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,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

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:

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 []