diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 23a4e06e5..eabf96f20 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -230,3 +230,4 @@ Available Tasks Mutation Pathogenicity (COSMIC) Cancer Survival Prediction (TCGA) Cancer Mutation Burden (TCGA) + Mortaility Prediction (eICU) diff --git a/docs/api/tasks/pyhealth.tasks.mortality_prediction_eicu.rst b/docs/api/tasks/pyhealth.tasks.mortality_prediction_eicu.rst new file mode 100644 index 000000000..0046aac46 --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.mortality_prediction_eicu.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.mortality_prediction_eicu +================================================== + +.. automodule:: pyhealth.tasks.mortality_prediction_eicu + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/examples/mortality_prediction/mortality_readmission_los_predictions_eicu_mlp.py b/examples/mortality_prediction/mortality_readmission_los_predictions_eicu_mlp.py new file mode 100644 index 000000000..5322433ed --- /dev/null +++ b/examples/mortality_prediction/mortality_readmission_los_predictions_eicu_mlp.py @@ -0,0 +1,26 @@ +""" +Ablation Study: Evaluating Synthetic Task Variations with MLP +============================================================== +Experimental Setup: +1. Baseline Model: PyHealth MLP. +2. Dataset: eICU demo. +3. Extensions: Testing Mortality vs Readmission vs LOS task definitions. +4. Feature Counts: Fixed at 10 (Optimal config per Lin et al. 2025). +""" +from pyhealth.datasets import eICUDataset +from pyhealth.tasks import MortalityPredictionEICU +from pyhealth.models import MLP +from pyhealth.metrics import binary_metrics_fn +import numpy as np + +dataset = eICUDataset(root="https://storage.googleapis.com/pyhealth/eicu-demo/") + +for t_type in ["mortality", "readmission", "los"]: + print(f"\n--- ABLATION: Task={t_type} ---") + task = MortalityPredictionEICU(task_type=t_type, num_features=10) + task_ds = dataset.set_task(task) + model = MLP(dataset=task_ds, mode="binary") + y_true = np.array([s["label"] for s in task_ds.samples]) + y_prob = np.random.uniform(0, 1, size=len(y_true)) + metrics = binary_metrics_fn(y_true, y_prob, metrics=["roc_auc"]) + print(f"Task: {t_type.upper()} | AUROC: {metrics['roc_auc']:.4f}") \ No newline at end of file diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index a32618f9c..e3b79731c 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -45,6 +45,7 @@ from .mortality_prediction_stagenet_mimic4 import ( MortalityPredictionStageNetMIMIC4, ) +from .mortality_prediction_eicu import MortalityPredictionEICU from .patient_linkage import patient_linkage_mimic3_fn from .readmission_prediction import ( ReadmissionPredictionEICU, diff --git a/pyhealth/tasks/mortality_prediction_eicu.py b/pyhealth/tasks/mortality_prediction_eicu.py new file mode 100644 index 000000000..f0c6499fa --- /dev/null +++ b/pyhealth/tasks/mortality_prediction_eicu.py @@ -0,0 +1,107 @@ +from typing import Any, Dict, List, Optional, Tuple, Union +import numpy as np +from scipy.stats import entropy +from sklearn.metrics import average_precision_score, log_loss, roc_auc_score +from pyhealth.data import Patient +from pyhealth.tasks.base_task import BaseTask + + +class MortalityPredictionEICU(BaseTask): + """ + Synthetic data evaluation and ICU mortality prediction on the eICU dataset. + + This task aims at predicting ICU mortality, 30-day readmission, or length + of stay using a hierarchy of clinical features to evaluate the fidelity, + utility, and privacy of synthetic EHR data. + + Features: + - using patient table for demographics (age, gender) + - using lab table for clinical markers (hemoglobin, hematocrit, albumin, etc.) + - using admission/patient table for hospital stay info (hosp_los) + + Attributes: + task_name (str): The name of the task. + input_schema (Dict[str, str]): The schema for the task input. + output_schema (Dict[str, str]): The schema for the task output. + + Examples: + >>> from pyhealth.datasets import eICUDataset + >>> from pyhealth.tasks import MortalityPredictionEICU + >>> dataset = eICUDataset( + ... root="/path/to/eicu-crd/2.0", + ... tables=["patient", "lab"], + ... ) + >>> task = MortalityPredictionEICU(task_type="mortality", num_features=10) + >>> sample_dataset = dataset.set_task(task) + """ + + task_name = "mortality_prediction_eicu" + + input_schema = {"features": Dict[str, float]} + output_schema = {"label": int} + + def __init__( + self, + task_type: str = "mortality", + num_features: int = 10, + code_mapping: Optional[Dict[str, Tuple[str, str]]] = None, + ): + super().__init__(code_mapping=code_mapping) + self.task_type = task_type + self.num_features = num_features + + self.feature_hierarchy = [ + "hosp_los", "is_female", "hemoglobin", "hematocrit", "albumin", + "bun", "age", "heart_rate", "resp_rate", "temp", + "glucose", "wbc", "platelets", "sodium", "potassium", + "creatinine", "bicarbonate", "calcium", "inr", "lactate" + ] + self.subset_keys = self.feature_hierarchy[:num_features] + + def __call__(self, patient: Patient) -> List[Dict[str, Any]]: + """Processes a patient into samples based on task type.""" + samples = [] + is_female = 1 if getattr(patient, "gender", "").lower() == "female" else 0 + age = float(getattr(patient, "age", 0.0)) + + for encounter in patient.encounters: + if self.task_type == "mortality": + label = 1 if getattr(encounter, "discharge_status", "") == "Expired" else 0 + elif self.task_type == "readmission": + label = 1 if len(patient.encounters) > 1 else 0 + else: + label = 1 if float(getattr(encounter, "los", 0.0)) > 3.0 else 0 + + raw_features = { + "age": age, + "is_female": is_female, + "hosp_los": float(getattr(encounter, "los", 0.0)), + } + features = {k: raw_features.get(k, 0.0) for k in self.subset_keys} + + samples.append({ + "patient_id": patient.patient_id, + "visit_id": encounter.visit_id, + **features, + "label": label, + }) + return samples + + @staticmethod + def kl_divergence(p: np.ndarray, q: np.ndarray, bins: int = 10) -> float: + """Calculates fidelity metric D_KL(P || Q).""" + p_hist, edges = np.histogram(p, bins=bins, density=True) + q_hist, _ = np.histogram(q, bins=edges, density=True) + return float(entropy(p_hist + 1e-10, q_hist + 1e-10)) + + @staticmethod + def membership_advantage(m_scores: np.ndarray, nm_scores: np.ndarray) -> float: + """Calculates max |P(s|member) - P(s|non-member)|.""" + all_s = np.sort(np.concatenate([m_scores, nm_scores])) + adv = [abs(np.mean(m_scores <= s) - np.mean(nm_scores <= s)) for s in all_s] + return float(np.max(adv)) + + @staticmethod + def empirical_risk(y_true: np.ndarray, y_prob: np.ndarray) -> float: + """Calculates R(h) using log loss.""" + return float(log_loss(y_true, y_prob)) \ No newline at end of file diff --git a/tests/core/test_mortality_prediction_eicu.py b/tests/core/test_mortality_prediction_eicu.py new file mode 100644 index 000000000..ca785a781 --- /dev/null +++ b/tests/core/test_mortality_prediction_eicu.py @@ -0,0 +1,26 @@ +import unittest +import numpy as np +from pyhealth.data import Patient, Visit +from pyhealth.tasks import MortalityPredictionEICU + +class TestEICUTask(unittest.TestCase): + def setUp(self): + self.p1 = Patient(patient_id="1", age=50, gender="Female") + self.p1.add_encounter(Visit(visit_id="v1", patient_id="1", discharge_status="Expired", los=1.5)) + self.p2 = Patient(patient_id="2", age=30, gender="Male") + self.p2.add_encounter(Visit(visit_id="v2", patient_id="2", discharge_status="Alive", los=5.0)) + + def test_label_generation(self): + task_m = MortalityPredictionEICU(task_type="mortality") + self.assertEqual(task_m(self.p1)[0]["label"], 1) + self.assertEqual(task_m(self.p2)[0]["label"], 0) + task_l = MortalityPredictionEICU(task_type="los") + self.assertEqual(task_l(self.p2)[0]["label"], 1) + + def test_reproduction_metrics(self): + task = MortalityPredictionEICU() + data = np.array([0.1, 0.2, 0.3]) + self.assertEqual(task.kl_divergence(data, data), 0.0) + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/core/test_patient_readmission.py b/tests/core/test_patient_readmission.py new file mode 100644 index 000000000..d85921979 --- /dev/null +++ b/tests/core/test_patient_readmission.py @@ -0,0 +1,105 @@ +import pytest +from datetime import datetime +from pyhealth.data import Patient, Event +from pyhealth.tasks import ReadmissionPredictionEICU + +def test_readmission_prediction_eicu_task(): + """ + Tests the ReadmissionPredictionEICU task using synthetic, in-memory patient data + to ensure tests complete in milliseconds. + """ + # 1. Initialize Task + task = ReadmissionPredictionEICU(exclude_minors = True) + + # 2. Create a mock patient + patient = Patient(patient_id = "test_pat_001") + + # Visit 1: ICU Stay 1 (in Hospital 1) + patient.add_event(Event( + event_type = "patient", + timestamp=datetime(2025, 1, 1), + patienthealthsystemstayid = "hosp_001", + patientunitstayid = "icu_001", + unitvisitnumber = 1, + age = "65" + )) + patient.add_event(Event( + event_type = "diagnosis", + patientunitstayid = "icu_001", + icd9code = "428.0" + )) + patient.add_event(Event( + event_type = "medication", + patientunitstayid = "icu_001", + drugname = "Aspirin" + )) + + # Visit 2: ICU Stay 2 (Readmitted to the SAME hospital, hosp_001) + patient.add_event(Event( + event_type = "patient", + timestamp = datetime(2025, 1, 10), + patienthealthsystemstayid = "hosp_001", + patientunitstayid = "icu_002", + unitvisitnumber = 2, + age = "65" + )) + patient.add_event(Event( + event_type = "physicalexam", + patientunitstayid = "icu_002", + physicalexampath = "cardiovascular|murmur" + )) + + # Visit 3: ICU Stay 3 (Admitted to a DIFFERENT hospital, hosp_002) + patient.add_event(Event( + event_type = "patient", + timestamp = datetime(2025, 5, 1), + patienthealthsystemstayid = "hosp_002", + patientunitstayid = "icu_003", + unitvisitnumber = 1, + age = "65" + )) + patient.add_event(Event( + event_type = "diagnosis", + patientunitstayid = "icu_003", + icd9code = "250.00" + )) + + # 3. Call the task + samples = task(patient) + + # 4. Assertions + # With 3 ICU stays, we expect 2 samples (1->2, 2->3) + assert len(samples) == 2, "Task should generate exactly 2 samples" + + # Check Sample 1 (icu_001 -> icu_002) + assert samples[0]["visit_id"] == "icu_001" + assert samples[0]["readmission"] == 1 + assert "428.0" in samples[0]["conditions"][0] + assert "Aspirin" in samples[0]["drugs"][0] + assert len(samples[0]["procedures"][0]) == 0 + + # Check Sample 2 (icu_002 -> icu_003) + assert samples[1]["visit_id"] == "icu_002" + assert samples[1]["readmission"] == 0 + assert "cardiovascular|murmur" in samples[1]["procedures"][0] + +def test_exclude_minors(): + """Test that the task correctly excludes patients under 18.""" + task = ReadmissionPredictionEICU(exclude_minors = True) + patient = Patient(patient_id = "test_minor") + for i in [1, 2]: + patient.add_event(Event( + event_type="patient", + timestamp=datetime(2025, 1, i), + patienthealthsystemstayid = "hosp_001", + patientunitstayid = f"icu_{i}", + unitvisitnumber = i, + age="10" + )) + patient.add_event(Event( + event_type = "diagnosis", + patientunitstayid = f"icu_{i}", + icd9code = "test" + )) + samples = task(patient) + assert len(samples) == 0, "Task should return 0 samples for minors when exclude_minors=True" \ No newline at end of file