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,154 @@
+"""
+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 standard PyHealth dataset → task → SampleDataset pipeline.
+
+Usage:
+    python mimic3_cf_circulatory_failure_logreg.py --root /path/to/mimic-iii
+"""
+
+import argparse
+
+import pandas as pd
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, recall_score, roc_auc_score
+from sklearn.model_selection import train_test_split
+
+from pyhealth.datasets import MIMIC3CirculatoryFailureDataset
+from pyhealth.tasks import CirculatoryFailurePredictionTask
+
+
+def samples_to_df(sample_dataset) -> pd.DataFrame:
+    """Converts a SampleDataset into a pandas DataFrame."""
+    rows = []
+    for i in range(len(sample_dataset)):
+        s = sample_dataset[i]
+        rows.append(
+            {
+                "patient_id": s["patient_id"],
+                "icustay_id": s["icustay_id"],
+                "gender": s.get("gender"),
+                "timestamp": s.get("timestamp"),
+                "map": to_scalar(s["map"]),
+                "map_diff": to_scalar(s["map_diff"]),
+                "label": int(to_scalar(s["label"])),
+            }
+        )
+    return pd.DataFrame(rows)
+
+
+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 to_scalar(x):
+    """Converts scalar tensor-like values to Python scalars."""
+    if hasattr(x, "item"):
+        return x.item()
+    return x
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="MIMIC-III circulatory failure prediction ablation study."
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        required=True,
+        help="Path to the unzipped MIMIC-III database directory.",
+    )
+    args = parser.parse_args()
+
+    dataset = MIMIC3CirculatoryFailureDataset(root=args.root)
+
+    # 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)
+        sample_dataset = dataset.set_task(task)
+        df = samples_to_df(sample_dataset)
+
+        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
+        advanced_metrics = evaluate_model(
+            df=df,
+            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[df["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()

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"

pyhealth/datasets/mimic3_cf.py

@@ -0,0 +1,91 @@
+"""
+MIMIC-III Circulatory Failure Dataset for PyHealth.
+
+Dataset:
+    MIMIC-III Clinical Database v1.4
+    https://physionet.org/content/mimiciii/1.4/
+
+Inspired by:
+    Hoche, M., Mineeva, O., Burger, M., Blasimme, A., & Ratsch, G. (2024). 
+    FAMEWS: A fairness auditing tool for medical early-warning systems. 
+    Proceedings of the Fifth Conference on Health, Inference, and Learning, 248, 297–311. PMLR. 
+    https://proceedings.mlr.press/v248/hoche24a.html
+
+Description:
+    Configures the MIMIC-III tables required for a circulatory-failure
+    early-warning task.  The dataset keeps data loading separate from
+    task logic; sample generation is handled by
+    ``CirculatoryFailurePredictionTask`` through the standard PyHealth
+    ``dataset.set_task(task)`` pipeline.
+
+Authors:
+    Kuang-Yu Wang ([email protected])
+    Ya Hsuan Yang ([email protected])
+"""
+
+import logging
+from pathlib import Path
+from typing import List, Optional
+from .base_dataset import BaseDataset
+
+logger = logging.getLogger(__name__)
+
+
+class MIMIC3CirculatoryFailureDataset(BaseDataset):
+    """MIMIC-III wrapper for circulatory failure early-warning prediction.
+
+    This dataset configures the MIMIC-III tables required for a
+    FAMEWS-inspired circulatory failure early-warning task. The dataset keeps
+    data loading separate from task logic; sample generation is handled by
+    ``CirculatoryFailurePredictionTask`` through the standard PyHealth
+    ``dataset.set_task(task)`` pipeline.
+
+    Args:
+        root: Root directory of the MIMIC-III dataset.
+        tables: Additional tables to load beyond the default tables.
+        dataset_name: Name of the dataset instance.
+        config_path: Path to the dataset config YAML file.
+        **kwargs: Additional keyword arguments passed to BaseDataset.
+
+    Examples:
+        >>> from pyhealth.datasets import MIMIC3CirculatoryFailureDataset
+        >>> from pyhealth.tasks import CirculatoryFailurePredictionTask
+        >>> dataset = MIMIC3CirculatoryFailureDataset(
+        ...     root="/path/to/mimic-iii",
+        ... )
+        >>> task = CirculatoryFailurePredictionTask(prediction_window_hours=12)
+        >>> sample_dataset = dataset.set_task(task)
+    """
+
+    def __init__(
+        self,
+        root: str,
+        tables: Optional[List[str]] = None,
+        dataset_name: Optional[str] = None,
+        config_path: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        """Initializes the MIMIC-III circulatory failure dataset."""
+        if config_path is None:
+            logger.info("No config path provided, using default config")
+            config_path = Path(__file__).parent / "configs" / "mimic3_cf.yaml"
+
+        default_tables = [
+            "patients",
+            "admissions",
+            "icustays",
+            "chartevents",
+        ]
+
+        if tables is None:
+            tables = default_tables
+        else:
+            tables = list(dict.fromkeys(default_tables + tables))
+
+        super().__init__(
+            root=root,
+            tables=tables,
+            dataset_name=dataset_name or "mimic3_cf",
+            config_path=str(config_path),
+            **kwargs,
+        )

pyhealth/tasks/__init__.py

@@ -67,3 +67,4 @@
     VariantClassificationClinVar,
 )
 from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task
+from .circulatory_failure_prediction import CirculatoryFailurePredictionTask