Fix post-outcome data leakage in StageNet mortality/LOS tasks

docs/api/tasks.rst

@@ -215,6 +215,7 @@ Available Tasks
     Drug Recommendation <tasks/pyhealth.tasks.drug_recommendation>
     EHR Generation <tasks/pyhealth.tasks.generate_ehr>
     Length of Stay Prediction <tasks/pyhealth.tasks.length_of_stay_prediction>
+    Length of Stay Prediction (StageNet MIMIC-IV) <tasks/pyhealth.tasks.length_of_stay_stagenet_mimic4>
     Medical Transcriptions Classification <tasks/pyhealth.tasks.MedicalTranscriptionsClassification>
     MPF Clinical Prediction (FHIR) <tasks/pyhealth.tasks.mpf_clinical_prediction>
     Mortality Prediction (Next Visit) <tasks/pyhealth.tasks.mortality_prediction>

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:

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

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

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: