From 3a7c630be5a7118f78255da6cf036a0ace7cceb9 Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 25 Aug 2026 16:58:06 -0500 Subject: [PATCH 1/2] Fix post-outcome data leakage in StageNet mortality/LOS tasks diagnoses_icd/procedures_icd events are timestamped at dischtime (per the MIMIC-IV config), so for the admission whose own outcome is being predicted, those codes are only known at-or-after the outcome. Labs were also pulled through discharge/death for that same admission. - MortalityPredictionStageNetMIMIC4: exclude diagnosis/procedure codes for the admission that ends in death, and cap that admission's labs to the first 48 hours after admission instead of through discharge. Earlier, already-resolved admissions are unaffected. - LengthOfStayStageNetMIMIC4: same treatment for the target admission (the one whose LOS is the label). Both changes only affect the single outcome-adjacent admission per patient; prior history is untouched. Reported metrics on these two tasks are expected to drop after this fix -- the previous numbers were inflated by the leak, not a legitimate baseline. Verified end-to-end against real MIMIC-IV demo data: terminal/target admission codes are excluded, late labs near the outcome are dropped, early labs and all prior-admission data are preserved, and patients whose only admission is the terminal/target one are correctly excluded (no leak-free signal available), matching the existing convention in dka.py. --- docs/api/tasks.rst | 1 + ...h.tasks.length_of_stay_stagenet_mimic4.rst | 7 + .../length_of_stay_mimic4_stagenet.py | 11 +- .../mortality_mimic4_stagenet_v2.py | 7 + .../tasks/length_of_stay_stagenet_mimic4.py | 121 ++++++++++++------ .../mortality_prediction_stagenet_mimic4.py | 110 +++++++++++----- .../test_stagenet_task_leakage_prevention.py | 99 ++++++++++++++ 7 files changed, 284 insertions(+), 72 deletions(-) create mode 100644 docs/api/tasks/pyhealth.tasks.length_of_stay_stagenet_mimic4.rst create mode 100644 tests/core/test_stagenet_task_leakage_prevention.py 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() From ef90261d46447e51d5b3da6ce8412eaa1d498444 Mon Sep 17 00:00:00 2001 From: lehendo Date: Sun, 30 Aug 2026 19:40:48 -0500 Subject: [PATCH 2/2] Fix asymmetric observation window in mortality leakage fix A reviewer (DarylOkeke) on PR #1205 pointed out that the leakage fix only restricted the death class: is_terminal_admission gated on hospital_expire_flag==1, so death cases had their terminal admission's ICD codes excluded and labs capped to 48h, while survivors' final admission kept full codes and labs through natural discharge. Two otherwise-identical patients would get systematically different feature richness based on the label itself -- a model could learn 'richer features -> survived' as a shortcut without any real clinical signal, and the existing test explicitly locked this asymmetry in as intended behavior. Fixed by determining a target admission the same way for both classes: the hospital_expire_flag==1 admission if the patient died, otherwise the chronologically last valid admission (mirroring how LengthOfStayStageNetMIMIC4 already treats its own target admission unconditionally). That admission's codes are excluded and its labs capped to TARGET_ADMISSION_INPUT_WINDOW_HOURS (renamed from TERMINAL_ADMISSION_INPUT_WINDOW_HOURS) regardless of the eventual label. Updated the existing survivor test, which asserted the old asymmetric behavior (all 3 admissions retained), to assert the corrected symmetric behavior instead, and added a direct cross-task symmetry check comparing mortality's and LOS's target-admission exclusion on the same survivor patient. Confirmed both new/updated tests fail against the pre-fix code (3 admissions instead of 2; leaked target-admission codes E1011/N179/ 5A1D70Z) and pass against the fix. --- .../mortality_mimic4_stagenet_v2.py | 12 +- .../mortality_prediction_stagenet_mimic4.py | 133 ++++++++++-------- .../test_stagenet_task_leakage_prevention.py | 57 ++++++-- 3 files changed, 133 insertions(+), 69 deletions(-) diff --git a/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py b/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py index 4ef989f33..e34399829 100644 --- a/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py +++ b/examples/mortality_prediction/mortality_mimic4_stagenet_v2.py @@ -9,10 +9,14 @@ 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 +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. """ diff --git a/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py b/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py index a86073d33..447444949 100644 --- a/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py +++ b/pyhealth/tasks/mortality_prediction_stagenet_mimic4.py @@ -26,16 +26,26 @@ class MortalityPredictionStageNetMIMIC4(BaseTask): 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 + ``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. - - 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. + - 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 @@ -63,11 +73,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 + # 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. @@ -149,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 = [] @@ -159,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: @@ -190,20 +220,9 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # Update previous admission time for next iteration previous_admission_time = admission_time - # 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 + is_target_admission = admission.hadm_id == target_hadm_id - if not is_terminal_admission: + if not is_target_admission: # Get diagnosis codes for this admission using hadm_id diagnoses_icd = patient.get_events( event_type="diagnoses_icd", @@ -233,14 +252,15 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: 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: + # 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.TERMINAL_ADMISSION_INPUT_WINDOW_HOURS), + + timedelta(hours=self.TARGET_ADMISSION_INPUT_WINDOW_HOURS), ) else: lab_window_end = admission_dischtime @@ -311,13 +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 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: + # 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 index 19955160f..ebb717bea 100644 --- a/tests/core/test_stagenet_task_leakage_prevention.py +++ b/tests/core/test_stagenet_task_leakage_prevention.py @@ -9,7 +9,7 @@ class TestStageNetTaskLeakagePrevention(unittest.TestCase): - """Regression tests for the terminal/target-admission leakage fix. + """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 @@ -17,7 +17,12 @@ class TestStageNetTaskLeakagePrevention(unittest.TestCase): 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. + 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 @@ -77,12 +82,18 @@ def test_los_excludes_target_admission_codes(self): 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. + 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) @@ -91,8 +102,34 @@ def test_mortality_survivor_unaffected(self): self.assertEqual(sample["mortality"], 0) _, icd_codes = sample["icd_codes"] - # All three admissions should contribute, none excluded. - self.assertEqual(len(icd_codes), 3) + # 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__":