Add MIMIC-III circulatory failure dataset and prediction task with ablation study

docs/api/datasets.rst

@@ -246,3 +246,4 @@ Available Datasets
     datasets/pyhealth.datasets.TCGAPRADDataset
     datasets/pyhealth.datasets.splitter
     datasets/pyhealth.datasets.utils
+    datasets/pyhealth.datasets.mimic3_cf

docs/api/datasets/pyhealth.datasets.mimic3_cf.rst

@@ -0,0 +1,26 @@
+pyhealth.datasets.mimic3_cf
+===========================
+
+Overview
+--------
+
+MIMIC3CirculatoryFailureDataset is a MIMIC-III based dataset for early warning
+prediction of circulatory failure.
+
+It constructs an ICU-stay-level cohort from PATIENTS, ADMISSIONS, and ICUSTAYS,
+and uses CHARTEVENTS to extract Mean Arterial Pressure (MAP) measurements.
+
+Circulatory failure is defined using a proxy event:
+
+- MAP < 65 mmHg
+
+For each ICU stay, the dataset identifies the first occurrence of this event and
+supports building task-ready patient records for downstream prediction tasks.
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.datasets.MIMIC3CirculatoryFailureDataset
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/tasks.rst

@@ -230,3 +230,4 @@ Available Tasks
     Mutation Pathogenicity (COSMIC) <tasks/pyhealth.tasks.MutationPathogenicityPrediction>
     Cancer Survival Prediction (TCGA) <tasks/pyhealth.tasks.CancerSurvivalPrediction>
     Cancer Mutation Burden (TCGA) <tasks/pyhealth.tasks.CancerMutationBurden>
+    Circulatory Failure Prediction <tasks/pyhealth.tasks.circulatory_failure_prediction>

docs/api/tasks/pyhealth.tasks.circulatory_failure_prediction.rst

@@ -0,0 +1,24 @@
+pyhealth.tasks.circulatory_failure_prediction
+=============================================
+
+Overview
+--------
+
+CirculatoryFailurePredictionTask defines a time-series prediction task for early
+detection of circulatory failure.
+
+The task predicts whether a patient will experience circulatory failure within
+the next 12 hours based on physiological measurements.
+
+Label definition:
+
+- label = 1 if circulatory failure occurs within the next 12 hours
+- label = 0 otherwise
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.tasks.CirculatoryFailurePredictionTask
+   :members:
+   :undoc-members:
+   :show-inheritance:

examples/mimic3_cf_circulatory_failure_logreg.py

@@ -0,0 +1,145 @@
+"""
+Example ablation script for MIMIC-III circulatory failure prediction.
+
+This script compares different prediction windows (6h, 12h, 24h) and
+feature settings using logistic regression. It is intended as an example
+usage script for the dataset-task pipeline and ablation study.
+"""
+
+from pyhealth.datasets import MIMIC3CirculatoryFailureDataset
+from pyhealth.tasks import CirculatoryFailurePredictionTask
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, roc_auc_score, recall_score
+
+
+def samples_to_df(samples: list[dict]) -> pd.DataFrame:
+    rows = []
+    for s in samples:
+        rows.append(
+            {
+                "patient_id": s["patient_id"],
+                "icustay_id": s["icustay_id"],
+                "gender": s["gender"],
+                "timestamp": s["timestamp"],
+                "map": s["features"]["map"],
+                "label": s["label"],
+            }
+        )
+    df = pd.DataFrame(rows)
+    return df
+
+
+def add_advanced_features(df: pd.DataFrame) -> pd.DataFrame:
+    """Add simple temporal features for the advanced setting."""
+    df = df.sort_values(["icustay_id", "timestamp"]).copy()
+    df["map_prev"] = df.groupby("icustay_id")["map"].shift(1)
+    df["map_diff"] = df["map"] - df["map_prev"]
+    df["map_prev"] = df["map_prev"].fillna(df["map"])
+    df["map_diff"] = df["map_diff"].fillna(0.0)
+    return df
+
+
+def evaluate_model(
+    df: pd.DataFrame,
+    feature_cols: list[str],
+    balanced: bool = False,
+) -> dict:
+    if df.empty or df["label"].nunique() < 2:
+        return {
+            "n_samples": len(df),
+            "accuracy": None,
+            "roc_auc": None,
+            "recall": None,
+        }
+
+    X = df[feature_cols]
+    y = df["label"]
+
+    X_train, X_test, y_train, y_test = train_test_split(
+        X,
+        y,
+        test_size=0.2,
+        random_state=42,
+        stratify=y,
+    )
+
+    model = LogisticRegression(
+        max_iter=1000,
+        class_weight="balanced" if balanced else None,
+    )
+    model.fit(X_train, y_train)
+
+    preds = model.predict(X_test)
+    probs = model.predict_proba(X_test)[:, 1]
+
+    return {
+        "n_samples": len(df),
+        "accuracy": accuracy_score(y_test, preds),
+        "roc_auc": roc_auc_score(y_test, probs),
+        "recall": recall_score(y_test, preds),
+    }
+
+
+def print_metrics(title: str, metrics: dict) -> None:
+    print(f"\n=== {title} ===")
+    print(f"n_samples: {metrics['n_samples']}")
+    print(f"accuracy: {metrics['accuracy']}")
+    print(f"roc_auc: {metrics['roc_auc']}")
+    print(f"recall: {metrics['recall']}")
+
+
+def main() -> None:
+    dataset = MIMIC3CirculatoryFailureDataset(
+        # path to the unzipped MIMIC-III database on your machine
+        root="mimic-iii-dataset"
+    )
+
+    # task ablation: prediction windows
+    for window in [6, 12, 24]:
+        print(f"\n############################")
+        print(f"Prediction window = {window}h")
+        print(f"############################")
+
+        task = CirculatoryFailurePredictionTask(prediction_window_hours=window)
+        samples = dataset.set_task(task, max_patients=100)
+        df = samples_to_df(samples)
+
+        print("\nSample preview:")
+        print(df.head())
+
+        # baseline setting
+        baseline_metrics = evaluate_model(
+            df=df,
+            feature_cols=["map"],
+            balanced=False,
+        )
+        print_metrics("Baseline: LogisticRegression(map)", baseline_metrics)
+
+        # advanced setting
+        df_adv = add_advanced_features(df)
+        advanced_metrics = evaluate_model(
+            df=df_adv,
+            feature_cols=["map", "map_diff"],
+            balanced=True,
+        )
+        print_metrics(
+            "Advanced: LogisticRegression(map + map_diff, balanced)",
+            advanced_metrics,
+        )
+
+        # subgroup fairness
+        for gender in ["M", "F"]:
+            subgroup_df = df_adv[df_adv["gender"] == gender].copy()
+            subgroup_metrics = evaluate_model(
+                df=subgroup_df,
+                feature_cols=["map", "map_diff"],
+                balanced=True,
+            )
+            print_metrics(f"Advanced subgroup gender={gender}", subgroup_metrics)
+
+
+if __name__ == "__main__":
+    main()

examples/mimic3_cf_example.py

@@ -0,0 +1,23 @@
+from pyhealth.datasets import MIMIC3CirculatoryFailureDataset
+from pyhealth.tasks import CirculatoryFailurePredictionTask
+
+
+def main():
+    dataset = MIMIC3CirculatoryFailureDataset(
+        root="/path/to/mimic3"
+    )
+
+    task = CirculatoryFailurePredictionTask(prediction_window_hours=12)
+
+    # apply task
+    samples = dataset.set_task(task, max_patients=5)
+
+    print(f"Total samples: {len(samples)}")
+
+    if samples:
+        print("Sample example:")
+        print(samples[0])
+
+
+if __name__ == "__main__":
+    main()

pyhealth/datasets/__init__.py

@@ -91,3 +91,4 @@ def __init__(self, *args, **kwargs):
     save_processors,
 )
 from .collate import collate_temporal
+from .mimic3_cf import MIMIC3CirculatoryFailureDataset

pyhealth/datasets/configs/mimic3_cf.yaml

@@ -0,0 +1,47 @@
+version: "1.4"
+tables:
+  patients:
+    file_path: "PATIENTS.csv.gz"
+    patient_id: "subject_id"
+    timestamp: null
+    attributes:
+      - "gender"
+      - "dob"
+      - "dod"
+      - "expire_flag"
+
+  admissions:
+    file_path: "ADMISSIONS.csv.gz"
+    patient_id: "subject_id"
+    timestamp: "admittime"
+    attributes:
+      - "hadm_id"
+      - "admittime"
+      - "dischtime"
+      - "deathtime"
+      - "hospital_expire_flag"
+      - "ethnicity"
+
+  icustays:
+    file_path: "ICUSTAYS.csv.gz"
+    patient_id: "subject_id"
+    timestamp: "intime"
+    attributes:
+      - "hadm_id"
+      - "icustay_id"
+      - "intime"
+      - "outtime"
+      - "first_careunit"
+      - "last_careunit"
+
+  chartevents:
+    file_path: "CHARTEVENTS.csv.gz"
+    patient_id: "subject_id"
+    timestamp: "charttime"
+    attributes:
+      - "hadm_id"
+      - "icustay_id"
+      - "itemid"
+      - "charttime"
+      - "value"
+      - "valuenum"