From 5ddd9a104949442d177b8f143c2d49c6e1fe269a Mon Sep 17 00:00:00 2001 From: Daryl Okeke Date: Sat, 5 Sep 2026 16:42:31 -0500 Subject: [PATCH] Fix multimodal mortality X-ray leakage --- .../pyhealth.tasks.mortality_prediction.rst | 7 +- .../multimodal_mimic4_minimal.py | 3 +- pyhealth/tasks/mortality_prediction.py | 21 ++- .../core/test_multimodal_mortality_leakage.py | 133 ++++++++++++++++++ 4 files changed, 158 insertions(+), 6 deletions(-) create mode 100644 tests/core/test_multimodal_mortality_leakage.py diff --git a/docs/api/tasks/pyhealth.tasks.mortality_prediction.rst b/docs/api/tasks/pyhealth.tasks.mortality_prediction.rst index 21aaca26f..81dbe1c1d 100644 --- a/docs/api/tasks/pyhealth.tasks.mortality_prediction.rst +++ b/docs/api/tasks/pyhealth.tasks.mortality_prediction.rst @@ -11,6 +11,11 @@ :undoc-members: :show-inheritance: +.. autoclass:: pyhealth.tasks.mortality_prediction.MultimodalMortalityPredictionMIMIC4 + :members: + :undoc-members: + :show-inheritance: + .. autoclass:: pyhealth.tasks.mortality_prediction.MortalityPredictionEICU :members: :undoc-members: @@ -24,4 +29,4 @@ .. autoclass:: pyhealth.tasks.mortality_prediction.MortalityPredictionOMOP :members: :undoc-members: - :show-inheritance: \ No newline at end of file + :show-inheritance: diff --git a/examples/mortality_prediction/multimodal_mimic4_minimal.py b/examples/mortality_prediction/multimodal_mimic4_minimal.py index 093e7764b..395729f50 100644 --- a/examples/mortality_prediction/multimodal_mimic4_minimal.py +++ b/examples/mortality_prediction/multimodal_mimic4_minimal.py @@ -24,7 +24,7 @@ num_workers=8 ) - # Apply multimodal task + # Apply multimodal task with X-rays available by prediction time task = MultimodalMortalityPredictionMIMIC4() samples = dataset.set_task(task, num_workers=8) @@ -32,4 +32,3 @@ sample = samples[0] print(sample) - diff --git a/pyhealth/tasks/mortality_prediction.py b/pyhealth/tasks/mortality_prediction.py index 249f717f3..3b67d0215 100644 --- a/pyhealth/tasks/mortality_prediction.py +++ b/pyhealth/tasks/mortality_prediction.py @@ -351,7 +351,7 @@ class MultimodalMortalityPredictionMIMIC4(BaseTask): Image Processing: - Uses image_path from MIMIC-CXR metadata directly - - Returns first available X-ray image path across all X-rays + - Uses X-rays available by the last included admission's discharge """ task_name: str = "MultimodalMortalityPredictionMIMIC4" @@ -575,6 +575,21 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if len(admissions_to_process) == 0: return [] + prediction_time = None + for admission in reversed(admissions_to_process): + try: + admission_dischtime = datetime.strptime( # noqa: DTZ007 + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + continue + if admission_dischtime >= admission.timestamp: + prediction_time = admission_dischtime + break + + if prediction_time is None: + return [] + # Get first admission time as reference for lab time calculations first_admission_time = admissions_to_process[0].timestamp @@ -591,8 +606,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # Get X-ray data (patient-level, not admission-specific) # Note: event types match table names in mimic4_cxr.yaml (negbio, metadata) - negbio_events = patient.get_events(event_type="negbio") - metadata_events = patient.get_events(event_type="metadata") + negbio_events = patient.get_events(event_type="negbio", end=prediction_time) + metadata_events = patient.get_events(event_type="metadata", end=prediction_time) # Process X-ray findings (aggregate across all X-rays) # NegBio findings attributes (from mimic4_cxr.yaml negbio table) diff --git a/tests/core/test_multimodal_mortality_leakage.py b/tests/core/test_multimodal_mortality_leakage.py new file mode 100644 index 000000000..a7dfc212d --- /dev/null +++ b/tests/core/test_multimodal_mortality_leakage.py @@ -0,0 +1,133 @@ +import unittest +from datetime import datetime, timedelta +from types import SimpleNamespace + +import polars as pl + +from pyhealth.tasks import MultimodalMortalityPredictionMIMIC4 + + +class DummyPatient: + def __init__(self, events): + self.patient_id = "patient" + self.events = events + + def get_events( + self, + event_type=None, + start=None, + end=None, + filters=None, + return_df=False, + ): + events = list(self.events.get(event_type, [])) + if start is not None: + events = [event for event in events if event.timestamp >= start] + if end is not None: + events = [event for event in events if event.timestamp <= end] + for field, operator, value in filters or []: + if operator == "==": + events = [ + event for event in events if getattr(event, field, None) == value + ] + + if return_df: + return pl.DataFrame( + { + "timestamp": [event.timestamp for event in events], + "labevents/itemid": [event.itemid for event in events], + "labevents/valuenum": [event.valuenum for event in events], + "labevents/storetime": [event.storetime for event in events], + } + ) + return events + + +def make_event(timestamp, **attributes): + return SimpleNamespace(timestamp=timestamp, **attributes) + + +def make_patient(*, include_available_xray): + admission_time = datetime(2025, 1, 1, 8) # noqa: DTZ001 + cutoff = admission_time + timedelta(days=1) + death_admission_time = cutoff + timedelta(days=3) + + metadata = [ + make_event( + death_admission_time, + image_path="future.jpg", + ) + ] + negbio = [ + make_event( + death_admission_time, + edema=1, + ) + ] + if include_available_xray: + metadata.insert(0, make_event(cutoff, image_path="available.jpg")) + negbio.insert(0, make_event(cutoff, cardiomegaly=1)) + + return DummyPatient( + { + "patients": [make_event(admission_time, anchor_age=50)], + "admissions": [ + make_event( + admission_time, + hadm_id="history", + dischtime=cutoff.strftime("%Y-%m-%d %H:%M:%S"), + hospital_expire_flag=0, + ), + make_event( + death_admission_time, + hadm_id="outcome", + dischtime=(death_admission_time + timedelta(days=1)).strftime( + "%Y-%m-%d %H:%M:%S" + ), + hospital_expire_flag=1, + ), + ], + "diagnoses_icd": [ + make_event(admission_time, hadm_id="history", icd_code="I10") + ], + "procedures_icd": [ + make_event(admission_time, hadm_id="history", icd_code="0W3P0ZZ") + ], + "prescriptions": [ + make_event(admission_time, hadm_id="history", ndc="0001") + ], + "discharge": [ + make_event(admission_time, hadm_id="history", text="Discharged") + ], + "radiology": [], + "labevents": [ + make_event( + admission_time + timedelta(hours=1), + itemid="50824", + valuenum=140.0, + storetime=(admission_time + timedelta(hours=1)).strftime( + "%Y-%m-%d %H:%M:%S" + ), + ) + ], + "metadata": metadata, + "negbio": negbio, + } + ) + + +class TestMultimodalMortalityLeakage(unittest.TestCase): + def test_excludes_future_xrays(self): + sample = MultimodalMortalityPredictionMIMIC4()( + make_patient(include_available_xray=True) + )[0] + + self.assertEqual(sample["image_path"], "available.jpg") + self.assertEqual(sample["negbio_findings"], ["cardiomegaly"]) + + def test_requires_xray_by_prediction_time(self): + samples = MultimodalMortalityPredictionMIMIC4()( + make_patient(include_available_xray=False) + ) + + self.assertEqual(samples, [])