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..e34399829 100644 --- a/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py +++ b/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py @@ -7,6 +7,17 @@ 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 target admission -- the one flagged as +the outcome if the patient died, otherwise the chronologically last +admission -- since discharge-timestamped codes are only known at-or-after +that admission's own outcome. That admission's labs are also capped to +the first TARGET_ADMISSION_INPUT_WINDOW_HOURS (default 48) hours after +admission instead of through discharge. This applies identically to both +death and survivor cases, so neither class gets a systematically richer +feature set than the other; 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..447444949 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,29 @@ 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 an + admission are only known at-or-after that admission's own + discharge/outcome. The *target* admission -- the one whose + ``hospital_expire_flag`` determines the label if the patient + died, or otherwise the chronologically last admission -- has its + codes excluded; 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 drawn late in the stay + (including, for a death, labs near the moment of death) are not + used as predictive features. Labs for earlier admissions are + unaffected. + - This restriction is applied identically regardless of the label: + a survivor's most recent admission is windowed the same way a + death case's terminal admission is. Restricting only the death + class (as an earlier version of this task did) would let a + model learn "richer features -> survived" as a shortcut, from + the asymmetric amount of information available per class, without + learning any real clinical signal. + Args: padding: Additional padding for StageNet processor to handle sequences longer than observed during training. Default: 0. @@ -50,6 +73,13 @@ class MortalityPredictionStageNetMIMIC4(BaseTask): task_name: str = "MortalityPredictionStageNetMIMIC4" + # For the target admission (see Data Leakage Prevention above), 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. Applied identically for + # both classes. + TARGET_ADMISSION_INPUT_WINDOW_HOURS: ClassVar[int] = 48 + def __init__(self, padding: int = 0): """Initialize task with optional padding parameter. @@ -130,6 +160,42 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if len(admissions) < 1: return [] + # Pre-parse admissions once to find valid ones and locate the + # target admission: the one flagged hospital_expire_flag==1 if the + # patient died, or otherwise the chronologically last valid + # admission. Its codes/labs get the leakage-safe restriction below, + # applied the same way regardless of the label, so both classes' + # target admission is windowed identically -- only its *content* + # differs, not the amount of information available. + valid_admissions = [] + target_hadm_id = None + for admission in admissions: + try: + admission_time = admission.timestamp + # MIMIC-IV timestamps carry no timezone; kept naive to match + # admission_time (also naive) for downstream comparisons. + admission_dischtime = datetime.strptime( # noqa: DTZ007 + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + continue + if admission_dischtime < admission_time: + continue + valid_admissions.append((admission, admission_time, admission_dischtime)) + if target_hadm_id is None: + try: + if int(admission.hospital_expire_flag) == 1: + target_hadm_id = admission.hadm_id + except (ValueError, TypeError, AttributeError): + pass + + if not valid_admissions: + return [] + + died = target_hadm_id is not None + if target_hadm_id is None: + target_hadm_id = valid_admissions[-1][0].hadm_id + # Initialize aggregated data structures # List of ICD codes (diagnoses + procedures) per visit all_icd_codes = [] @@ -140,25 +206,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # Track previous admission timestamp for interval calculation previous_admission_time = None - # Track if patient had any mortality event - final_mortality = 0 - # Process each admission - for i, admission in enumerate(admissions): - # Parse admission and discharge times - try: - admission_time = admission.timestamp - admission_dischtime = datetime.strptime( - admission.dischtime, "%Y-%m-%d %H:%M:%S" - ) - except (ValueError, AttributeError): - # Skip if timestamps invalid - continue - - # Skip if discharge is before admission (data quality issue) - if admission_dischtime < admission_time: - continue - + for admission, admission_time, admission_dischtime in valid_admissions: # Calculate time from previous admission (in hours) # First admission will have time = 0 if previous_admission_time is None: @@ -171,47 +220,55 @@ 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 - try: - if int(admission.hospital_expire_flag) == 1: - final_mortality = 1 - 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 + is_target_admission = admission.hadm_id == target_hadm_id - if visit_icd_codes: - all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) + if not is_target_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 target admission, + # cap the window to the first TARGET_ADMISSION_INPUT_WINDOW_HOURS + # hours after admission instead of through discharge, regardless + # of whether this patient's outcome is death or survival. + if is_target_admission: + lab_window_end = min( + admission_dischtime, + admission_time + + timedelta(hours=self.TARGET_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 +285,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 +331,16 @@ 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 target admission: for a death, any further + # admissions would be chronologically impossible for this + # patient, but we guard against including their data in case + # of a data-quality inconsistency; for a survivor, this is the + # chronologically last admission anyway. + if is_target_admission: + break + + final_mortality = 1 if died else 0 + # 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..ebb717bea --- /dev/null +++ b/tests/core/test_stagenet_task_leakage_prevention.py @@ -0,0 +1,136 @@ +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 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 -- and, for mortality specifically, that this + restriction applies identically to death and survivor cases. An + earlier version of the fix restricted only the death class, which let + a model learn "richer features -> survived" as a shortcut from the + asymmetric amount of information available per class, rather than any + real clinical signal. + """ + + @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_target_admission_also_excluded(self): + """Regression test for the class-asymmetry leak: a survivor's most + recent admission must be windowed the same way a death case's + terminal admission is, not left fully unrestricted. + + Patient 10001 has three admissions (19999, 20001, then 20002), all + survived -- the same fixture used by the LOS test above, which + already (correctly) excludes 20002 as its target admission. Before + this fix, MortalityPredictionStageNetMIMIC4 only restricted the + death class, so this same patient's 20002 codes would leak in here + while being correctly excluded for LOS -- an inconsistency that is + itself evidence of the asymmetry. + """ + 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"] + # Only the two non-target (19999, 20001) admissions should + # contribute code lists; the target (20002) must not. + 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 survivor features", + ) + + def test_mortality_and_los_treat_same_survivor_identically(self): + """Direct symmetry check: for the same survivor, mortality and LOS + must agree on which admission is the target and exclude the exact + same code set from it -- confirming the mortality task no longer + gives survivors a privileged, unrestricted view of their own most + recent admission relative to what LOS already does correctly. + """ + patient = self.dataset.get_patient("10001") + mortality_codes = MortalityPredictionStageNetMIMIC4()(patient)[0]["icd_codes"][1] + los_codes = LengthOfStayStageNetMIMIC4()(patient)[0]["icd_codes"][1] + + mortality_flat = sorted(code for visit in mortality_codes for code in visit) + los_flat = sorted(code for visit in los_codes for code in visit) + self.assertEqual(mortality_flat, los_flat) + + +if __name__ == "__main__": + unittest.main()