diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index bdaa9599a..4f887ea4f 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -215,6 +215,7 @@ Available Tasks Drug Recommendation EHR Generation Length of Stay Prediction + Length of Stay Prediction (StageNet MIMIC-IV) Medical Transcriptions Classification MPF Clinical Prediction (FHIR) Mortality Prediction (Next Visit) diff --git a/docs/api/tasks/pyhealth.tasks.length_of_stay_stagenet_mimic4.rst b/docs/api/tasks/pyhealth.tasks.length_of_stay_stagenet_mimic4.rst new file mode 100644 index 000000000..02d7c8e1f --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.length_of_stay_stagenet_mimic4.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.length_of_stay_stagenet_mimic4 +=============================================== + +.. autoclass:: pyhealth.tasks.length_of_stay_stagenet_mimic4.LengthOfStayStageNetMIMIC4 + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/length_of_stay/length_of_stay_mimic4_stagenet.py b/examples/length_of_stay/length_of_stay_mimic4_stagenet.py index 3c560fe5b..2c6b52c96 100644 --- a/examples/length_of_stay/length_of_stay_mimic4_stagenet.py +++ b/examples/length_of_stay/length_of_stay_mimic4_stagenet.py @@ -1,12 +1,19 @@ """ -Example of using StageNet for mortality prediction on MIMIC-IV. +Example of using StageNet for length of stay prediction on MIMIC-IV. This example demonstrates: 1. Loading MIMIC-IV data -2. Applying the MortalityPredictionStageNetMIMIC4 task +2. Applying the LengthOfStayStageNetMIMIC4 task 3. Creating a SampleDataset with StageNet processors 4. Training a StageNet model 5. Testing with synthetic hold-out set (unseen codes, varying lengths) + +Note: to prevent leakage, LengthOfStayStageNetMIMIC4 excludes the +diagnosis/procedure codes of the target admission (the one whose LOS is +the label) since they're only known at-or-after its own discharge, and +caps that admission's labs to the first TARGET_ADMISSION_INPUT_WINDOW_HOURS +(default 48) hours after admission instead of through discharge. Earlier, +already-resolved admissions are unaffected. """ import os diff --git a/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py b/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py index acf9598d0..4ef989f33 100644 --- a/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py +++ b/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py @@ -7,6 +7,13 @@ 3. Creating a SampleDataset with StageNet processors 4. Training a StageNet model 5. Testing with synthetic hold-out set (unseen codes, varying lengths) + +Note: to prevent leakage, MortalityPredictionStageNetMIMIC4 excludes the +diagnosis/procedure codes of the admission that ends in death (they're only +known at-or-after discharge) and caps that admission's labs to the first +TERMINAL_ADMISSION_INPUT_WINDOW_HOURS (default 48) hours after admission +instead of through discharge. Earlier, already-resolved admissions are +unaffected. """ import os diff --git a/pyhealth/tasks/length_of_stay_stagenet_mimic4.py b/pyhealth/tasks/length_of_stay_stagenet_mimic4.py index be05a22b6..f8efe3225 100644 --- a/pyhealth/tasks/length_of_stay_stagenet_mimic4.py +++ b/pyhealth/tasks/length_of_stay_stagenet_mimic4.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, ClassVar, Dict, List, Tuple import polars as pl @@ -26,6 +26,20 @@ class LengthOfStayStageNetMIMIC4(BaseTask): - 10D vectors, one value per lab category (first observed per category per timestamp, missing -> None) + Data Leakage Prevention + ------------------------ + - The prediction target is the LOS of the most recent (target) + admission. ``diagnoses_icd``/``procedures_icd`` events are timestamped + at ``dischtime`` (per the MIMIC-IV config) -- i.e. at-or-after that + admission's own discharge, which is what determines its LOS label. + Those codes are excluded for the target admission; codes from earlier, + already-resolved admissions are unaffected. + - Labs for the target admission are restricted to the first + ``TARGET_ADMISSION_INPUT_WINDOW_HOURS`` hours after admission, rather + than through discharge, so labs that are only available because the + stay ran long are not used to predict its own length. Labs for + earlier admissions are unaffected. + Args: padding: Optional padding forwarded to the StageNet processor for nested sequences. Default is 0. @@ -43,6 +57,11 @@ class LengthOfStayStageNetMIMIC4(BaseTask): task_name: str = "LengthOfStayStageNetMIMIC4" + # For the target admission (whose LOS is the label), only labs drawn + # within this many hours of admission are used as features, rather than + # the full window through discharge. + TARGET_ADMISSION_INPUT_WINDOW_HOURS: ClassVar[int] = 48 + def __init__(self, padding: int = 0): self.padding = padding self.input_schema: Dict[str, Tuple[str, Dict[str, Any]]] = { @@ -100,14 +119,11 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if len(admissions) < 1: return [] - all_icd_codes: List[List[str]] = [] - all_icd_times: List[float] = [] - all_lab_values: List[List[Any]] = [] - all_lab_times: List[float] = [] - - previous_admission_time = None - target_los_category = None - + # Parse and validate admission times once. This also lets us + # identify the target admission (the most recent valid one, whose + # LOS is the label) as the last entry, so its data can be + # restricted below without a second parsing pass. + valid_admissions = [] for admission in admissions: try: admission_time = admission.timestamp @@ -116,13 +132,28 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) except (ValueError, AttributeError): continue - if discharge_time < admission_time: continue + valid_admissions.append((admission, admission_time, discharge_time)) + + if not valid_admissions: + return [] + + target_hadm_id = valid_admissions[-1][0].hadm_id + + all_icd_codes: list[list[str]] = [] + all_icd_times: list[float] = [] + all_lab_values: list[list[Any]] = [] + all_lab_times: list[float] = [] + previous_admission_time = None + target_los_category = None + + for admission, admission_time, discharge_time in valid_admissions: # Label from the most recent valid admission encountered los_days = (discharge_time - admission_time).days target_los_category = categorize_los(los_days) + is_target_admission = admission.hadm_id == target_hadm_id if previous_admission_time is None: time_from_previous = 0.0 @@ -133,36 +164,52 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: previous_admission_time = admission_time - diagnoses_icd = patient.get_events( - event_type="diagnoses_icd", - filters=[("hadm_id", "==", admission.hadm_id)], - ) - visit_diagnoses = [ - event.icd_code - for event in diagnoses_icd - if hasattr(event, "icd_code") and event.icd_code - ] - - procedures_icd = patient.get_events( - event_type="procedures_icd", - filters=[("hadm_id", "==", admission.hadm_id)], - ) - visit_procedures = [ - event.icd_code - for event in procedures_icd - if hasattr(event, "icd_code") and event.icd_code - ] - - visit_icd_codes = visit_diagnoses + visit_procedures - - if visit_icd_codes: - all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) + # Diagnoses/procedures are timestamped at dischtime, so for the + # target admission they're only known at-or-after its own LOS + # outcome -- exclude them. Earlier admissions are unaffected. + if not is_target_admission: + diagnoses_icd = patient.get_events( + event_type="diagnoses_icd", + filters=[("hadm_id", "==", admission.hadm_id)], + ) + visit_diagnoses = [ + event.icd_code + for event in diagnoses_icd + if hasattr(event, "icd_code") and event.icd_code + ] + + procedures_icd = patient.get_events( + event_type="procedures_icd", + filters=[("hadm_id", "==", admission.hadm_id)], + ) + visit_procedures = [ + event.icd_code + for event in procedures_icd + if hasattr(event, "icd_code") and event.icd_code + ] + + visit_icd_codes = visit_diagnoses + visit_procedures + + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + + # For the target admission, cap the lab window to the first + # TARGET_ADMISSION_INPUT_WINDOW_HOURS hours after admission + # instead of through discharge. + if is_target_admission: + lab_window_end = min( + discharge_time, + admission_time + + timedelta(hours=self.TARGET_ADMISSION_INPUT_WINDOW_HOURS), + ) + else: + lab_window_end = discharge_time labevents_df = patient.get_events( event_type="labevents", start=admission_time, - end=discharge_time, + end=lab_window_end, return_df=True, ) @@ -177,7 +224,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) ) labevents_df = labevents_df.filter( - pl.col("labevents/storetime") <= discharge_time + pl.col("labevents/storetime") <= lab_window_end ) if labevents_df.height > 0: diff --git a/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py b/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py index 4c0505f2d..a86073d33 100644 --- a/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py +++ b/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, ClassVar, Dict, List, Tuple import polars as pl @@ -24,6 +24,19 @@ class MortalityPredictionStageNetMIMIC4(BaseTask): - Multiple itemids per category → take first observed value - Missing categories → None/NaN in vector + Data Leakage Prevention: + - ``diagnoses_icd``/``procedures_icd`` events are timestamped at + ``dischtime`` (per the MIMIC-IV config), so codes recorded for the + admission that ends in death are only known at-or-after the + outcome. Those codes are excluded for the terminal (mortality) + admission; codes from earlier, already-resolved admissions are + unaffected. + - Labs for the terminal admission are restricted to the first + ``TERMINAL_ADMISSION_INPUT_WINDOW_HOURS`` hours after admission, + rather than through discharge, so death-adjacent labs drawn near + the moment of death are not used as predictive features. Labs for + earlier admissions are unaffected. + Args: padding: Additional padding for StageNet processor to handle sequences longer than observed during training. Default: 0. @@ -50,6 +63,12 @@ class MortalityPredictionStageNetMIMIC4(BaseTask): task_name: str = "MortalityPredictionStageNetMIMIC4" + # For the admission that ends in death, only labs drawn within this many + # hours of admission are used as features (mirrors the fixed prediction + # window used by InHospitalMortalityMIMIC4), rather than the full window + # through discharge/death. + TERMINAL_ADMISSION_INPUT_WINDOW_HOURS: ClassVar[int] = 48 + def __init__(self, padding: int = 0): """Initialize task with optional padding parameter. @@ -171,47 +190,65 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # Update previous admission time for next iteration previous_admission_time = admission_time - # Update mortality label if this admission had mortality + # Determine if this admission is the terminal (mortality) one. + # diagnoses_icd/procedures_icd are timestamped at dischtime, so + # codes for this admission are only known at-or-after the + # outcome and must be excluded; labs are capped to an early + # fixed window instead of through discharge/death. + is_terminal_admission = False try: if int(admission.hospital_expire_flag) == 1: final_mortality = 1 + is_terminal_admission = True except (ValueError, TypeError, AttributeError): pass - # Get diagnosis codes for this admission using hadm_id - diagnoses_icd = patient.get_events( - event_type="diagnoses_icd", - filters=[("hadm_id", "==", admission.hadm_id)], - ) - visit_diagnoses = [ - event.icd_code - for event in diagnoses_icd - if hasattr(event, "icd_code") and event.icd_code - ] - - # Get procedure codes for this admission using hadm_id - procedures_icd = patient.get_events( - event_type="procedures_icd", - filters=[("hadm_id", "==", admission.hadm_id)], - ) - visit_procedures = [ - event.icd_code - for event in procedures_icd - if hasattr(event, "icd_code") and event.icd_code - ] - - # Combine diagnoses and procedures into single ICD code list - visit_icd_codes = visit_diagnoses + visit_procedures - - if visit_icd_codes: - all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) + if not is_terminal_admission: + # Get diagnosis codes for this admission using hadm_id + diagnoses_icd = patient.get_events( + event_type="diagnoses_icd", + filters=[("hadm_id", "==", admission.hadm_id)], + ) + visit_diagnoses = [ + event.icd_code + for event in diagnoses_icd + if hasattr(event, "icd_code") and event.icd_code + ] + + # Get procedure codes for this admission using hadm_id + procedures_icd = patient.get_events( + event_type="procedures_icd", + filters=[("hadm_id", "==", admission.hadm_id)], + ) + visit_procedures = [ + event.icd_code + for event in procedures_icd + if hasattr(event, "icd_code") and event.icd_code + ] + + # Combine diagnoses and procedures into single ICD code list + visit_icd_codes = visit_diagnoses + visit_procedures + + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + + # Get lab events for this admission. For the terminal admission, + # cap the window to the first TERMINAL_ADMISSION_INPUT_WINDOW_HOURS + # hours after admission instead of through discharge/death. + if is_terminal_admission: + lab_window_end = min( + admission_dischtime, + admission_time + + timedelta(hours=self.TERMINAL_ADMISSION_INPUT_WINDOW_HOURS), + ) + else: + lab_window_end = admission_dischtime - # Get lab events for this admission labevents_df = patient.get_events( event_type="labevents", start=admission_time, - end=admission_dischtime, + end=lab_window_end, return_df=True, ) @@ -228,7 +265,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) ) labevents_df = labevents_df.filter( - pl.col("labevents/storetime") <= admission_dischtime + pl.col("labevents/storetime") <= lab_window_end ) if labevents_df.height > 0: @@ -274,6 +311,13 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_values.append(lab_vector) all_lab_times.append(time_from_admission) + # Stop after the terminal admission: any further admissions + # would be chronologically impossible for this patient, but we + # guard against including their data in case of a data-quality + # inconsistency. + if is_terminal_admission: + break + # Skip if no lab events (required for this task) if len(all_lab_values) == 0: return [] diff --git a/tests/core/test_stagenet_task_leakage_prevention.py b/tests/core/test_stagenet_task_leakage_prevention.py new file mode 100644 index 000000000..19955160f --- /dev/null +++ b/tests/core/test_stagenet_task_leakage_prevention.py @@ -0,0 +1,99 @@ +import unittest +from pathlib import Path + +from pyhealth.datasets import MIMIC4Dataset +from pyhealth.tasks.length_of_stay_stagenet_mimic4 import LengthOfStayStageNetMIMIC4 +from pyhealth.tasks.mortality_prediction_stagenet_mimic4 import ( + MortalityPredictionStageNetMIMIC4, +) + + +class TestStageNetTaskLeakagePrevention(unittest.TestCase): + """Regression tests for the terminal/target-admission leakage fix. + + ``diagnoses_icd``/``procedures_icd`` events are timestamped at + ``dischtime`` (see pyhealth/datasets/configs/mimic4_ehr.yaml), so codes + recorded for the admission whose own outcome (mortality or LOS) is being + predicted are only known at-or-after that outcome. These tests verify + that MortalityPredictionStageNetMIMIC4 and LengthOfStayStageNetMIMIC4 + exclude that admission's codes while leaving earlier, already-resolved + admissions unaffected. + """ + + @classmethod + def setUpClass(cls): + test_dir = Path(__file__).parent.parent.parent + root = str(test_dir / "test-resources" / "core" / "mimic4demo") + tables = ["diagnoses_icd", "procedures_icd", "prescriptions", "labevents"] + cls.dataset = MIMIC4Dataset(ehr_root=root, ehr_tables=tables) + + def test_mortality_excludes_terminal_admission_codes(self): + """Patient 10003 has two admissions (20005, then terminal 20006). + + Codes from the non-terminal admission (20005) must be present; + codes from the terminal admission (20006), which are only known at + its own dischtime, must not leak into the features. + """ + patient = self.dataset.get_patient("10003") + samples = MortalityPredictionStageNetMIMIC4()(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertEqual(sample["mortality"], 1) + + _, icd_codes = sample["icd_codes"] + flat_codes = [code for visit in icd_codes for code in visit] + + for code in ["E1010", "I10", "5A1955Z"]: + self.assertIn(code, flat_codes) + for code in ["E1011", "N170", "I509", "5A1D70Z", "02HV33Z"]: + self.assertNotIn( + code, + flat_codes, + f"terminal-admission code {code} leaked into features", + ) + + def test_los_excludes_target_admission_codes(self): + """Patient 10001 has three admissions (19999, 20001, then 20002), + all survived. The LOS label comes from the most recent (target) + admission (20002), whose codes must be excluded; codes unique to + the two earlier admissions must still be present. + """ + patient = self.dataset.get_patient("10001") + samples = LengthOfStayStageNetMIMIC4()(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + + _, icd_codes = sample["icd_codes"] + # Only the two non-target admissions should contribute code lists. + self.assertEqual(len(icd_codes), 2) + + flat_codes = [code for visit in icd_codes for code in visit] + for code in ["E1010", "E1165", "I10", "5A1955Z", "3E0G76Z"]: + self.assertIn(code, flat_codes) + for code in ["E1011", "N179", "5A1D70Z"]: + self.assertNotIn( + code, + flat_codes, + f"target-admission code {code} leaked into LOS features", + ) + + def test_mortality_survivor_unaffected(self): + """A patient with no terminal admission keeps full historical data. + + This is a regression check: the leakage fix must only change + behavior for the terminal/target admission, not for patients who + never trigger it. + """ + patient = self.dataset.get_patient("10001") + samples = MortalityPredictionStageNetMIMIC4()(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertEqual(sample["mortality"], 0) + + _, icd_codes = sample["icd_codes"] + # All three admissions should contribute, none excluded. + self.assertEqual(len(icd_codes), 3) + + +if __name__ == "__main__": + unittest.main()