From 6eda1a857ff4f1c9374b3da4e3b6ac81cd18b1c0 Mon Sep 17 00:00:00 2001 From: John Wu Date: Mon, 8 Sep 2025 15:14:29 -0500 Subject: [PATCH 1/6] update to ckd --- pyhealth/tasks/ckd_surv.py | 467 +++++++++++++++++++++++++++++++++++++ 1 file changed, 467 insertions(+) create mode 100644 pyhealth/tasks/ckd_surv.py diff --git a/pyhealth/tasks/ckd_surv.py b/pyhealth/tasks/ckd_surv.py new file mode 100644 index 000000000..ff57d1047 --- /dev/null +++ b/pyhealth/tasks/ckd_surv.py @@ -0,0 +1,467 @@ +from datetime import datetime +from typing import Any, Dict, List, Optional, Literal +import polars as pl + +from .base_task import BaseTask + + +class MIMIC4CKDSurvAnalysis(BaseTask): + """CKD survival analysis task with simplified configuration. + + eGFR calculation methodology adapted from: + - Original implementation: pkgs.data.utils.calculate_eGFR() + - Formula source: pkgs.data.store.get_egfr_df() + - Reference: CKD-EPI 2021 formula (https://pubmed.ncbi.nlm.nih.gov/34554658/) + """ + + # Private class variables for settings + _SURVIVAL_SETTINGS = ["time_invariant", "time_variant", "heterogeneous"] + _CKD_CODES = ["N183", "N184", "N185", "585.3", "585.4", "585.5"] + _ESRD_CODES = ["N186", "Z992", "585.6", "V42.0"] + _CREATININE_ITEMIDS = ["50912", "52546"] + _PROTEIN_ITEMIDS = ["50976"] + _ALBUMIN_ITEMIDS = ["50862"] + + # Gender constants (using MIMIC-IV native string values) + _MALE_GENDER = "M" # Male patients + _FEMALE_GENDER = "F" # Female patients + + # CKD-EPI 2021 equation constants (from pkgs.data.utils.calculate_eGFR) + _BASE_COEFFICIENT = 141 # Original uses 142 in utils.py + _AGE_FACTOR = 0.993 # Annual age decline factor + _FEMALE_ADJUSTMENT = 1.018 # Female gender boost factor + + # Gender-specific creatinine thresholds and exponents + _MALE_CREAT_THRESHOLD = 0.9 # mg/dL + _FEMALE_CREAT_THRESHOLD = 0.7 # mg/dL + _MALE_ALPHA_EXPONENT = -0.411 # For creatinine ≤ 0.9 + _FEMALE_ALPHA_EXPONENT = -0.329 # For creatinine ≤ 0.7 + _BETA_EXPONENT = -1.209 # For creatinine > threshold (both genders) + + def __init__( + self, + setting: Literal[ + "time_invariant", "time_variant", "heterogeneous" + ] = "time_invariant", + min_age: int = 18, + prediction_window_days: int = 365 * 5, + ): + + if setting not in self._SURVIVAL_SETTINGS: + raise ValueError(f"Setting must be one of {self._SURVIVAL_SETTINGS}") + + self.setting = setting + self.min_age = min_age + self.prediction_window_days = prediction_window_days + self.task_name = f"MIMIC4CKDSurvAnalysis_{self.setting}" + self.input_schema, self.output_schema = self._configure_schemas() + + def _configure_schemas(self) -> tuple[Dict[str, str], Dict[str, str]]: + """Configure schemas based on survival setting.""" + base_input = {"demographics": "List[str]", "age": "float", "gender": "str"} + + base_output = {"duration_days": "float", "has_esrd": "int"} + + if self.setting == "time_invariant": + base_input.update({"baseline_egfr": "float", "comorbidities": "List[str]"}) + elif self.setting == "time_variant": + base_input.update( + { + "lab_measurements": "List[Dict[str, Any]]" # [{"timestamp": days, "egfr": value}, ...] + } + ) + else: # heterogeneous + base_input.update( + { + "lab_measurements": "List[Dict[str, Any]]", # [{"timestamp": days, "egfr": value, "protein": value, "missing_egfr": bool}, ...] + "missing_indicators": "List[str]", + } + ) + + return base_input, base_output + + def filter_patients(self, df: pl.LazyFrame) -> pl.LazyFrame: + """Filter for CKD patients with required lab data.""" + ckd_patients = ( + df.filter(pl.col("event_type") == "diagnoses_icd") + .filter(pl.col("diagnoses_icd/icd_code").is_in(self._CKD_CODES)) + .select("patient_id") + .unique() + ) + + lab_patients = ( + df.filter(pl.col("event_type") == "labevents") + .filter(pl.col("labevents/itemid").is_in(self._CREATININE_ITEMIDS)) + .select("patient_id") + .unique() + ) + + valid_patients = ckd_patients.join(lab_patients, on="patient_id", how="inner") + return df.filter( + pl.col("patient_id").is_in(valid_patients.select("patient_id")) + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + """Process patient for survival analysis.""" + # Get demographics + demographics = patient.get_events(event_type="patients") + if not demographics: + return [] + + demo = demographics[0] + age = int(demo.anchor_age or 0) + gender = (demo.gender or "").upper() + + if gender not in [self._MALE_GENDER, self._FEMALE_GENDER]: + return [] # Skip patients with invalid/missing gender + + if age < self.min_age: + return [] + + # Get CKD baseline date + ckd_diagnoses = patient.get_events(event_type="diagnoses_icd") + ckd_events = [e for e in ckd_diagnoses if e.icd_code in self._CKD_CODES] + + if not ckd_events: + return [] + + baseline_date = min(e.timestamp for e in ckd_events) + + # Get ESRD outcome + esrd_events = [e for e in ckd_diagnoses if e.icd_code in self._ESRD_CODES] + + if esrd_events: + esrd_date = min( + e.timestamp for e in esrd_events if e.timestamp > baseline_date + ) + has_esrd = 1 + duration_days = (esrd_date - baseline_date).days + else: + has_esrd = 0 + # Get all events to find last observation - FIXED: Filter out None timestamps + all_events = patient.get_events() # Get all events + # Filter out events with None timestamps before finding max + valid_events = [e for e in all_events if e.timestamp is not None] + + if valid_events: + last_event = max(valid_events, key=lambda x: x.timestamp) + duration_days = min( + (last_event.timestamp - baseline_date).days, + self.prediction_window_days, + ) + else: + # Fallback: use prediction window if no valid timestamps found + duration_days = self.prediction_window_days + + if duration_days <= 0: + return [] + + # Process by setting + if self.setting == "time_invariant": + return self._process_time_invariant( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + elif self.setting == "time_variant": + return self._process_time_variant( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + else: # heterogeneous + return self._process_heterogeneous( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + + def _process_time_invariant( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for time-invariant analysis.""" + lab_events = patient.get_events(event_type="labevents") + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ] + + if not creatinine_events: + return [] + + # Validate and find baseline creatinine + valid_creatinine_events = [] + for e in creatinine_events: + try: + creatinine_value = float(e.valuenum) + if creatinine_value > 0: # Must be positive + valid_creatinine_events.append((e, creatinine_value)) + except (ValueError, TypeError): + continue # Skip non-numeric values + + if not valid_creatinine_events: + return [] + + # Find baseline creatinine closest to baseline_date + baseline_event, baseline_creatinine_value = min( + valid_creatinine_events, + key=lambda x: abs((x[0].timestamp - baseline_date).days), + ) + + egfr = self._calculate_egfr(baseline_creatinine_value, age, gender) + + # Get comorbidities + diagnoses = patient.get_events(event_type="diagnoses_icd") + comorbidities = [ + e.icd_code for e in diagnoses if e.timestamp <= baseline_date and e.icd_code + ] + + # Get race + admissions = patient.get_events(event_type="admissions") + race = admissions[0].race if admissions else "unknown" + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str, race], + "baseline_egfr": egfr, + "comorbidities": comorbidities, + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _process_time_variant( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for time-varying analysis.""" + lab_events = patient.get_events(event_type="labevents") + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ] + + if len(creatinine_events) < 2: + return [] + + # Sort by time and create labeled measurements + creatinine_events.sort(key=lambda x: x.timestamp) + lab_measurements = [] + + for e in creatinine_events: + # Validate and convert creatinine value + try: + creatinine_value = float(e.valuenum) + if creatinine_value <= 0: + continue # Skip invalid values + except (ValueError, TypeError): + continue # Skip non-numeric values + + days_from_baseline = (e.timestamp - baseline_date).days + egfr_value = self._calculate_egfr(creatinine_value, age, gender) + + lab_measurements.append( + { + "timestamp": days_from_baseline, + "egfr": egfr_value, + "creatinine": creatinine_value, + } + ) + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str], + "lab_measurements": lab_measurements, + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _process_heterogeneous( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for heterogeneous analysis with missing indicators.""" + lab_events = patient.get_events(event_type="labevents") + + # Get multiple biomarkers with validation + def validate_and_convert_lab_events(events, itemids): + """Helper to validate and convert lab values.""" + valid_events = [] + for e in events: + if e.itemid in itemids and e.valuenum is not None: + try: + value = float(e.valuenum) + if value > 0: # Must be positive for lab values + valid_events.append((e, value)) + except (ValueError, TypeError): + continue + return valid_events + + creatinine_events = validate_and_convert_lab_events( + lab_events, self._CREATININE_ITEMIDS + ) + protein_events = validate_and_convert_lab_events( + lab_events, self._PROTEIN_ITEMIDS + ) + albumin_events = validate_and_convert_lab_events( + lab_events, self._ALBUMIN_ITEMIDS + ) + + if not creatinine_events: + return [] + + # Create time-aligned measurements with all biomarkers + measurements_by_time = {} + + # Add creatinine/eGFR measurements + for e, creatinine_value in creatinine_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + egfr = self._calculate_egfr(creatinine_value, age, gender) + + if days not in measurements_by_time: + measurements_by_time[days] = {"timestamp": days} + + measurements_by_time[days].update( + { + "egfr": egfr, + "creatinine": creatinine_value, + "missing_egfr": False, + } + ) + + # Add protein measurements + for e, protein_value in protein_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + + if days not in measurements_by_time: + measurements_by_time[days] = { + "timestamp": days, + "missing_egfr": True, + } + + measurements_by_time[days].update( + {"protein": protein_value, "missing_protein": False} + ) + + # Add albumin measurements + for e, albumin_value in albumin_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + + if days not in measurements_by_time: + measurements_by_time[days] = { + "timestamp": days, + "missing_egfr": True, + } + + measurements_by_time[days].update( + {"albumin": albumin_value, "missing_albumin": False} + ) + + if len(measurements_by_time) < 2: + return [] + + # Convert to sorted list and fill missing indicators + lab_measurements = [] + for days in sorted(measurements_by_time.keys()): + measurement = measurements_by_time[days] + + # Set missing indicators for features not present + measurement.setdefault("missing_egfr", True) + measurement.setdefault("missing_protein", True) + measurement.setdefault("missing_albumin", True) + + # Set default values for missing features + measurement.setdefault("egfr", 0.0) + measurement.setdefault("protein", 0.0) + measurement.setdefault("albumin", 0.0) + measurement.setdefault("creatinine", 0.0) + + lab_measurements.append(measurement) + + # Collect all missing indicator types present + missing_indicators = set() + for measurement in lab_measurements: + for key, value in measurement.items(): + if key.startswith("missing_") and value: + missing_indicators.add(key) + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str], + "lab_measurements": lab_measurements, + "missing_indicators": list(missing_indicators), + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _calculate_egfr(self, creatinine: float, age: int, gender: str) -> float: + """Calculate eGFR using simplified CKD-EPI equation. + + Implementation adapted from original MIMIC-IV analysis code: + - Source file: pkgs.data.utils.calculate_eGFR() + - Formula: CKD-EPI 2021 (https://pubmed.ncbi.nlm.nih.gov/34554658/) + - Original coefficient: 142 (updated from 141 in this implementation) + + CKD-EPI Formula Constants (from original utils.py): + - 0.9/0.7: Gender-specific creatinine thresholds (mg/dL) + - 0.993: Age factor per year + - 1.018: Female gender adjustment factor + - -0.411/-0.329: Alpha exponents for creatinine ≤ threshold + - -1.209: Beta exponent for creatinine > threshold (both genders) + + Args: + creatinine: Serum creatinine in mg/dL (MIMIC-IV native units) + age: Patient age in years + gender: Gender string ('M' for male, 'F' for female) + + Returns: + Estimated GFR in mL/min/1.73m² + """ + # Validate inputs (following original validation) + if creatinine <= 0: + raise ValueError(f"Invalid creatinine value: {creatinine}") + if gender not in [self._MALE_GENDER, self._FEMALE_GENDER]: + raise ValueError(f"Invalid gender: {gender}") + + # Ensure creatinine is float for calculations + creatinine = float(creatinine) + + if gender == self._MALE_GENDER: # Male + return ( + self._BASE_COEFFICIENT + * min(creatinine / self._MALE_CREAT_THRESHOLD, 1) + ** self._MALE_ALPHA_EXPONENT + * max(creatinine / self._MALE_CREAT_THRESHOLD, 1) ** self._BETA_EXPONENT + * self._AGE_FACTOR**age + ) + else: # Female (gender == self._FEMALE_GENDER) + return ( + self._BASE_COEFFICIENT + * min(creatinine / self._FEMALE_CREAT_THRESHOLD, 1) + ** self._FEMALE_ALPHA_EXPONENT + * max(creatinine / self._FEMALE_CREAT_THRESHOLD, 1) + ** self._BETA_EXPONENT + * self._AGE_FACTOR**age + * self._FEMALE_ADJUSTMENT + ) From 25ac35ab90fa0cbfb8c8831598b7561c2360b1d2 Mon Sep 17 00:00:00 2001 From: John Wu Date: Mon, 8 Sep 2025 15:14:29 -0500 Subject: [PATCH 2/6] update to ckd --- pyhealth/tasks/ckd_surv.py | 467 +++++++++++++++++++++++++++++++++++++ 1 file changed, 467 insertions(+) create mode 100644 pyhealth/tasks/ckd_surv.py diff --git a/pyhealth/tasks/ckd_surv.py b/pyhealth/tasks/ckd_surv.py new file mode 100644 index 000000000..ff57d1047 --- /dev/null +++ b/pyhealth/tasks/ckd_surv.py @@ -0,0 +1,467 @@ +from datetime import datetime +from typing import Any, Dict, List, Optional, Literal +import polars as pl + +from .base_task import BaseTask + + +class MIMIC4CKDSurvAnalysis(BaseTask): + """CKD survival analysis task with simplified configuration. + + eGFR calculation methodology adapted from: + - Original implementation: pkgs.data.utils.calculate_eGFR() + - Formula source: pkgs.data.store.get_egfr_df() + - Reference: CKD-EPI 2021 formula (https://pubmed.ncbi.nlm.nih.gov/34554658/) + """ + + # Private class variables for settings + _SURVIVAL_SETTINGS = ["time_invariant", "time_variant", "heterogeneous"] + _CKD_CODES = ["N183", "N184", "N185", "585.3", "585.4", "585.5"] + _ESRD_CODES = ["N186", "Z992", "585.6", "V42.0"] + _CREATININE_ITEMIDS = ["50912", "52546"] + _PROTEIN_ITEMIDS = ["50976"] + _ALBUMIN_ITEMIDS = ["50862"] + + # Gender constants (using MIMIC-IV native string values) + _MALE_GENDER = "M" # Male patients + _FEMALE_GENDER = "F" # Female patients + + # CKD-EPI 2021 equation constants (from pkgs.data.utils.calculate_eGFR) + _BASE_COEFFICIENT = 141 # Original uses 142 in utils.py + _AGE_FACTOR = 0.993 # Annual age decline factor + _FEMALE_ADJUSTMENT = 1.018 # Female gender boost factor + + # Gender-specific creatinine thresholds and exponents + _MALE_CREAT_THRESHOLD = 0.9 # mg/dL + _FEMALE_CREAT_THRESHOLD = 0.7 # mg/dL + _MALE_ALPHA_EXPONENT = -0.411 # For creatinine ≤ 0.9 + _FEMALE_ALPHA_EXPONENT = -0.329 # For creatinine ≤ 0.7 + _BETA_EXPONENT = -1.209 # For creatinine > threshold (both genders) + + def __init__( + self, + setting: Literal[ + "time_invariant", "time_variant", "heterogeneous" + ] = "time_invariant", + min_age: int = 18, + prediction_window_days: int = 365 * 5, + ): + + if setting not in self._SURVIVAL_SETTINGS: + raise ValueError(f"Setting must be one of {self._SURVIVAL_SETTINGS}") + + self.setting = setting + self.min_age = min_age + self.prediction_window_days = prediction_window_days + self.task_name = f"MIMIC4CKDSurvAnalysis_{self.setting}" + self.input_schema, self.output_schema = self._configure_schemas() + + def _configure_schemas(self) -> tuple[Dict[str, str], Dict[str, str]]: + """Configure schemas based on survival setting.""" + base_input = {"demographics": "List[str]", "age": "float", "gender": "str"} + + base_output = {"duration_days": "float", "has_esrd": "int"} + + if self.setting == "time_invariant": + base_input.update({"baseline_egfr": "float", "comorbidities": "List[str]"}) + elif self.setting == "time_variant": + base_input.update( + { + "lab_measurements": "List[Dict[str, Any]]" # [{"timestamp": days, "egfr": value}, ...] + } + ) + else: # heterogeneous + base_input.update( + { + "lab_measurements": "List[Dict[str, Any]]", # [{"timestamp": days, "egfr": value, "protein": value, "missing_egfr": bool}, ...] + "missing_indicators": "List[str]", + } + ) + + return base_input, base_output + + def filter_patients(self, df: pl.LazyFrame) -> pl.LazyFrame: + """Filter for CKD patients with required lab data.""" + ckd_patients = ( + df.filter(pl.col("event_type") == "diagnoses_icd") + .filter(pl.col("diagnoses_icd/icd_code").is_in(self._CKD_CODES)) + .select("patient_id") + .unique() + ) + + lab_patients = ( + df.filter(pl.col("event_type") == "labevents") + .filter(pl.col("labevents/itemid").is_in(self._CREATININE_ITEMIDS)) + .select("patient_id") + .unique() + ) + + valid_patients = ckd_patients.join(lab_patients, on="patient_id", how="inner") + return df.filter( + pl.col("patient_id").is_in(valid_patients.select("patient_id")) + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + """Process patient for survival analysis.""" + # Get demographics + demographics = patient.get_events(event_type="patients") + if not demographics: + return [] + + demo = demographics[0] + age = int(demo.anchor_age or 0) + gender = (demo.gender or "").upper() + + if gender not in [self._MALE_GENDER, self._FEMALE_GENDER]: + return [] # Skip patients with invalid/missing gender + + if age < self.min_age: + return [] + + # Get CKD baseline date + ckd_diagnoses = patient.get_events(event_type="diagnoses_icd") + ckd_events = [e for e in ckd_diagnoses if e.icd_code in self._CKD_CODES] + + if not ckd_events: + return [] + + baseline_date = min(e.timestamp for e in ckd_events) + + # Get ESRD outcome + esrd_events = [e for e in ckd_diagnoses if e.icd_code in self._ESRD_CODES] + + if esrd_events: + esrd_date = min( + e.timestamp for e in esrd_events if e.timestamp > baseline_date + ) + has_esrd = 1 + duration_days = (esrd_date - baseline_date).days + else: + has_esrd = 0 + # Get all events to find last observation - FIXED: Filter out None timestamps + all_events = patient.get_events() # Get all events + # Filter out events with None timestamps before finding max + valid_events = [e for e in all_events if e.timestamp is not None] + + if valid_events: + last_event = max(valid_events, key=lambda x: x.timestamp) + duration_days = min( + (last_event.timestamp - baseline_date).days, + self.prediction_window_days, + ) + else: + # Fallback: use prediction window if no valid timestamps found + duration_days = self.prediction_window_days + + if duration_days <= 0: + return [] + + # Process by setting + if self.setting == "time_invariant": + return self._process_time_invariant( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + elif self.setting == "time_variant": + return self._process_time_variant( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + else: # heterogeneous + return self._process_heterogeneous( + patient, baseline_date, age, gender, duration_days, has_esrd + ) + + def _process_time_invariant( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for time-invariant analysis.""" + lab_events = patient.get_events(event_type="labevents") + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ] + + if not creatinine_events: + return [] + + # Validate and find baseline creatinine + valid_creatinine_events = [] + for e in creatinine_events: + try: + creatinine_value = float(e.valuenum) + if creatinine_value > 0: # Must be positive + valid_creatinine_events.append((e, creatinine_value)) + except (ValueError, TypeError): + continue # Skip non-numeric values + + if not valid_creatinine_events: + return [] + + # Find baseline creatinine closest to baseline_date + baseline_event, baseline_creatinine_value = min( + valid_creatinine_events, + key=lambda x: abs((x[0].timestamp - baseline_date).days), + ) + + egfr = self._calculate_egfr(baseline_creatinine_value, age, gender) + + # Get comorbidities + diagnoses = patient.get_events(event_type="diagnoses_icd") + comorbidities = [ + e.icd_code for e in diagnoses if e.timestamp <= baseline_date and e.icd_code + ] + + # Get race + admissions = patient.get_events(event_type="admissions") + race = admissions[0].race if admissions else "unknown" + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str, race], + "baseline_egfr": egfr, + "comorbidities": comorbidities, + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _process_time_variant( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for time-varying analysis.""" + lab_events = patient.get_events(event_type="labevents") + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ] + + if len(creatinine_events) < 2: + return [] + + # Sort by time and create labeled measurements + creatinine_events.sort(key=lambda x: x.timestamp) + lab_measurements = [] + + for e in creatinine_events: + # Validate and convert creatinine value + try: + creatinine_value = float(e.valuenum) + if creatinine_value <= 0: + continue # Skip invalid values + except (ValueError, TypeError): + continue # Skip non-numeric values + + days_from_baseline = (e.timestamp - baseline_date).days + egfr_value = self._calculate_egfr(creatinine_value, age, gender) + + lab_measurements.append( + { + "timestamp": days_from_baseline, + "egfr": egfr_value, + "creatinine": creatinine_value, + } + ) + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str], + "lab_measurements": lab_measurements, + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _process_heterogeneous( + self, patient, baseline_date, age, gender, duration_days, has_esrd + ): + """Process for heterogeneous analysis with missing indicators.""" + lab_events = patient.get_events(event_type="labevents") + + # Get multiple biomarkers with validation + def validate_and_convert_lab_events(events, itemids): + """Helper to validate and convert lab values.""" + valid_events = [] + for e in events: + if e.itemid in itemids and e.valuenum is not None: + try: + value = float(e.valuenum) + if value > 0: # Must be positive for lab values + valid_events.append((e, value)) + except (ValueError, TypeError): + continue + return valid_events + + creatinine_events = validate_and_convert_lab_events( + lab_events, self._CREATININE_ITEMIDS + ) + protein_events = validate_and_convert_lab_events( + lab_events, self._PROTEIN_ITEMIDS + ) + albumin_events = validate_and_convert_lab_events( + lab_events, self._ALBUMIN_ITEMIDS + ) + + if not creatinine_events: + return [] + + # Create time-aligned measurements with all biomarkers + measurements_by_time = {} + + # Add creatinine/eGFR measurements + for e, creatinine_value in creatinine_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + egfr = self._calculate_egfr(creatinine_value, age, gender) + + if days not in measurements_by_time: + measurements_by_time[days] = {"timestamp": days} + + measurements_by_time[days].update( + { + "egfr": egfr, + "creatinine": creatinine_value, + "missing_egfr": False, + } + ) + + # Add protein measurements + for e, protein_value in protein_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + + if days not in measurements_by_time: + measurements_by_time[days] = { + "timestamp": days, + "missing_egfr": True, + } + + measurements_by_time[days].update( + {"protein": protein_value, "missing_protein": False} + ) + + # Add albumin measurements + for e, albumin_value in albumin_events: + if e.timestamp >= baseline_date: + days = (e.timestamp - baseline_date).days + + if days not in measurements_by_time: + measurements_by_time[days] = { + "timestamp": days, + "missing_egfr": True, + } + + measurements_by_time[days].update( + {"albumin": albumin_value, "missing_albumin": False} + ) + + if len(measurements_by_time) < 2: + return [] + + # Convert to sorted list and fill missing indicators + lab_measurements = [] + for days in sorted(measurements_by_time.keys()): + measurement = measurements_by_time[days] + + # Set missing indicators for features not present + measurement.setdefault("missing_egfr", True) + measurement.setdefault("missing_protein", True) + measurement.setdefault("missing_albumin", True) + + # Set default values for missing features + measurement.setdefault("egfr", 0.0) + measurement.setdefault("protein", 0.0) + measurement.setdefault("albumin", 0.0) + measurement.setdefault("creatinine", 0.0) + + lab_measurements.append(measurement) + + # Collect all missing indicator types present + missing_indicators = set() + for measurement in lab_measurements: + for key, value in measurement.items(): + if key.startswith("missing_") and value: + missing_indicators.add(key) + + age_group = "elderly" if age >= 65 else "adult" + gender_str = "male" if gender == self._MALE_GENDER else "female" + + sample = { + "patient_id": patient.patient_id, + "demographics": [age_group, gender_str], + "lab_measurements": lab_measurements, + "missing_indicators": list(missing_indicators), + "age": float(age), + "gender": gender, + "duration_days": float(duration_days), + "has_esrd": has_esrd, + } + + return [sample] # Single sample wrapped in list for consistent interface + + def _calculate_egfr(self, creatinine: float, age: int, gender: str) -> float: + """Calculate eGFR using simplified CKD-EPI equation. + + Implementation adapted from original MIMIC-IV analysis code: + - Source file: pkgs.data.utils.calculate_eGFR() + - Formula: CKD-EPI 2021 (https://pubmed.ncbi.nlm.nih.gov/34554658/) + - Original coefficient: 142 (updated from 141 in this implementation) + + CKD-EPI Formula Constants (from original utils.py): + - 0.9/0.7: Gender-specific creatinine thresholds (mg/dL) + - 0.993: Age factor per year + - 1.018: Female gender adjustment factor + - -0.411/-0.329: Alpha exponents for creatinine ≤ threshold + - -1.209: Beta exponent for creatinine > threshold (both genders) + + Args: + creatinine: Serum creatinine in mg/dL (MIMIC-IV native units) + age: Patient age in years + gender: Gender string ('M' for male, 'F' for female) + + Returns: + Estimated GFR in mL/min/1.73m² + """ + # Validate inputs (following original validation) + if creatinine <= 0: + raise ValueError(f"Invalid creatinine value: {creatinine}") + if gender not in [self._MALE_GENDER, self._FEMALE_GENDER]: + raise ValueError(f"Invalid gender: {gender}") + + # Ensure creatinine is float for calculations + creatinine = float(creatinine) + + if gender == self._MALE_GENDER: # Male + return ( + self._BASE_COEFFICIENT + * min(creatinine / self._MALE_CREAT_THRESHOLD, 1) + ** self._MALE_ALPHA_EXPONENT + * max(creatinine / self._MALE_CREAT_THRESHOLD, 1) ** self._BETA_EXPONENT + * self._AGE_FACTOR**age + ) + else: # Female (gender == self._FEMALE_GENDER) + return ( + self._BASE_COEFFICIENT + * min(creatinine / self._FEMALE_CREAT_THRESHOLD, 1) + ** self._FEMALE_ALPHA_EXPONENT + * max(creatinine / self._FEMALE_CREAT_THRESHOLD, 1) + ** self._BETA_EXPONENT + * self._AGE_FACTOR**age + * self._FEMALE_ADJUSTMENT + ) From 969954fc06a9b1fb40783202ef9d507a8694e3ae Mon Sep 17 00:00:00 2001 From: John Wu Date: Thu, 18 Sep 2025 19:22:19 -0500 Subject: [PATCH 3/6] Test cases, and processors alignments --- pyhealth/tasks/ckd_surv.py | 177 +++++++++++++++-------------- tests/core/test_mimic4_ckd_surv.py | 121 ++++++++++++++++++++ 2 files changed, 211 insertions(+), 87 deletions(-) create mode 100644 tests/core/test_mimic4_ckd_surv.py diff --git a/pyhealth/tasks/ckd_surv.py b/pyhealth/tasks/ckd_surv.py index ff57d1047..1680977c2 100644 --- a/pyhealth/tasks/ckd_surv.py +++ b/pyhealth/tasks/ckd_surv.py @@ -1,17 +1,22 @@ -from datetime import datetime -from typing import Any, Dict, List, Optional, Literal +from typing import Any, Dict, List, Literal, Union, Type import polars as pl from .base_task import BaseTask +from pyhealth.processors import ( + SequenceProcessor, + TensorProcessor, + RawProcessor, +) class MIMIC4CKDSurvAnalysis(BaseTask): """CKD survival analysis task with simplified configuration. - eGFR calculation methodology adapted from: + eGFR calculation methodology adapted from: - Original implementation: pkgs.data.utils.calculate_eGFR() - Formula source: pkgs.data.store.get_egfr_df() - - Reference: CKD-EPI 2021 formula (https://pubmed.ncbi.nlm.nih.gov/34554658/) + - Reference: CKD-EPI 2021 formula + (https://pubmed.ncbi.nlm.nih.gov/34554658/) """ # Private class variables for settings @@ -56,25 +61,43 @@ def __init__( self.task_name = f"MIMIC4CKDSurvAnalysis_{self.setting}" self.input_schema, self.output_schema = self._configure_schemas() - def _configure_schemas(self) -> tuple[Dict[str, str], Dict[str, str]]: - """Configure schemas based on survival setting.""" - base_input = {"demographics": "List[str]", "age": "float", "gender": "str"} + def _configure_schemas( + self, + ) -> tuple[Dict[str, Union[str, Type]], Dict[str, Union[str, Type]]]: + """Configure schemas based on survival setting. - base_output = {"duration_days": "float", "has_esrd": "int"} + Use registered processors: + - "sequence" for categorical lists + (e.g., demographics, gender, comorbidities) + - "tensor" for numeric values (e.g., age, eGFR, durations, labels) + """ + base_input: Dict[str, Union[str, Type]] = { + "demographics": SequenceProcessor, + "age": TensorProcessor, + "gender": SequenceProcessor, + } + + base_output: Dict[str, Union[str, Type]] = { + "duration_days": TensorProcessor, + "has_esrd": TensorProcessor, + } if self.setting == "time_invariant": - base_input.update({"baseline_egfr": "float", "comorbidities": "List[str]"}) - elif self.setting == "time_variant": base_input.update( { - "lab_measurements": "List[Dict[str, Any]]" # [{"timestamp": days, "egfr": value}, ...] + "baseline_egfr": TensorProcessor, + "comorbidities": SequenceProcessor, } ) + elif self.setting == "time_variant": + # Use raw processor for time series list of dicts + base_input.update({"lab_measurements": RawProcessor}) else: # heterogeneous + # Raw lab measurements with sequence for missing indicators base_input.update( { - "lab_measurements": "List[Dict[str, Any]]", # [{"timestamp": days, "egfr": value, "protein": value, "missing_egfr": bool}, ...] - "missing_indicators": "List[str]", + "lab_measurements": RawProcessor, + "missing_indicators": SequenceProcessor, } ) @@ -138,7 +161,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: duration_days = (esrd_date - baseline_date).days else: has_esrd = 0 - # Get all events to find last observation - FIXED: Filter out None timestamps + # Get all events to find last observation + # FIXED: filter out None timestamps all_events = patient.get_events() # Get all events # Filter out events with None timestamps before finding max valid_events = [e for e in all_events if e.timestamp is not None] @@ -178,9 +202,11 @@ def _process_time_invariant( creatinine_events = [ e for e in lab_events - if e.itemid in self._CREATININE_ITEMIDS - and e.valuenum is not None - and e.timestamp >= baseline_date + if ( + e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ) ] if not creatinine_events: @@ -191,29 +217,29 @@ def _process_time_invariant( for e in creatinine_events: try: creatinine_value = float(e.valuenum) - if creatinine_value > 0: # Must be positive + if creatinine_value > 0: valid_creatinine_events.append((e, creatinine_value)) except (ValueError, TypeError): - continue # Skip non-numeric values + continue if not valid_creatinine_events: return [] - # Find baseline creatinine closest to baseline_date - baseline_event, baseline_creatinine_value = min( + # Closest to baseline + _, baseline_creatinine_value = min( valid_creatinine_events, key=lambda x: abs((x[0].timestamp - baseline_date).days), ) egfr = self._calculate_egfr(baseline_creatinine_value, age, gender) - # Get comorbidities + # Comorbidities before baseline diagnoses = patient.get_events(event_type="diagnoses_icd") comorbidities = [ e.icd_code for e in diagnoses if e.timestamp <= baseline_date and e.icd_code ] - # Get race + # Race from admissions admissions = patient.get_events(event_type="admissions") race = admissions[0].race if admissions else "unknown" @@ -226,12 +252,11 @@ def _process_time_invariant( "baseline_egfr": egfr, "comorbidities": comorbidities, "age": float(age), - "gender": gender, + "gender": [gender], "duration_days": float(duration_days), "has_esrd": has_esrd, } - - return [sample] # Single sample wrapped in list for consistent interface + return [sample] def _process_time_variant( self, patient, baseline_date, age, gender, duration_days, has_esrd @@ -241,30 +266,30 @@ def _process_time_variant( creatinine_events = [ e for e in lab_events - if e.itemid in self._CREATININE_ITEMIDS - and e.valuenum is not None - and e.timestamp >= baseline_date + if ( + e.itemid in self._CREATININE_ITEMIDS + and e.valuenum is not None + and e.timestamp >= baseline_date + ) ] if len(creatinine_events) < 2: return [] - # Sort by time and create labeled measurements + # Chronological order creatinine_events.sort(key=lambda x: x.timestamp) lab_measurements = [] for e in creatinine_events: - # Validate and convert creatinine value try: creatinine_value = float(e.valuenum) if creatinine_value <= 0: - continue # Skip invalid values + continue except (ValueError, TypeError): - continue # Skip non-numeric values + continue days_from_baseline = (e.timestamp - baseline_date).days egfr_value = self._calculate_egfr(creatinine_value, age, gender) - lab_measurements.append( { "timestamp": days_from_baseline, @@ -281,12 +306,11 @@ def _process_time_variant( "demographics": [age_group, gender_str], "lab_measurements": lab_measurements, "age": float(age), - "gender": gender, + "gender": [gender], "duration_days": float(duration_days), "has_esrd": has_esrd, } - - return [sample] # Single sample wrapped in list for consistent interface + return [sample] def _process_heterogeneous( self, patient, baseline_date, age, gender, duration_days, has_esrd @@ -294,45 +318,35 @@ def _process_heterogeneous( """Process for heterogeneous analysis with missing indicators.""" lab_events = patient.get_events(event_type="labevents") - # Get multiple biomarkers with validation - def validate_and_convert_lab_events(events, itemids): - """Helper to validate and convert lab values.""" - valid_events = [] + # Validator for lab values + def validate(events, itemids): + out = [] for e in events: if e.itemid in itemids and e.valuenum is not None: try: - value = float(e.valuenum) - if value > 0: # Must be positive for lab values - valid_events.append((e, value)) + v = float(e.valuenum) + if v > 0: + out.append((e, v)) except (ValueError, TypeError): continue - return valid_events + return out - creatinine_events = validate_and_convert_lab_events( - lab_events, self._CREATININE_ITEMIDS - ) - protein_events = validate_and_convert_lab_events( - lab_events, self._PROTEIN_ITEMIDS - ) - albumin_events = validate_and_convert_lab_events( - lab_events, self._ALBUMIN_ITEMIDS - ) + creatinine_events = validate(lab_events, self._CREATININE_ITEMIDS) + protein_events = validate(lab_events, self._PROTEIN_ITEMIDS) + albumin_events = validate(lab_events, self._ALBUMIN_ITEMIDS) if not creatinine_events: return [] - # Create time-aligned measurements with all biomarkers - measurements_by_time = {} + measurements_by_time: Dict[int, Dict[str, Any]] = {} - # Add creatinine/eGFR measurements + # Creatinine/eGFR for e, creatinine_value in creatinine_events: if e.timestamp >= baseline_date: days = (e.timestamp - baseline_date).days egfr = self._calculate_egfr(creatinine_value, age, gender) - if days not in measurements_by_time: measurements_by_time[days] = {"timestamp": days} - measurements_by_time[days].update( { "egfr": egfr, @@ -341,32 +355,28 @@ def validate_and_convert_lab_events(events, itemids): } ) - # Add protein measurements + # Protein for e, protein_value in protein_events: if e.timestamp >= baseline_date: days = (e.timestamp - baseline_date).days - if days not in measurements_by_time: measurements_by_time[days] = { "timestamp": days, "missing_egfr": True, } - measurements_by_time[days].update( {"protein": protein_value, "missing_protein": False} ) - # Add albumin measurements + # Albumin for e, albumin_value in albumin_events: if e.timestamp >= baseline_date: days = (e.timestamp - baseline_date).days - if days not in measurements_by_time: measurements_by_time[days] = { "timestamp": days, "missing_egfr": True, } - measurements_by_time[days].update( {"albumin": albumin_value, "missing_albumin": False} ) @@ -374,28 +384,22 @@ def validate_and_convert_lab_events(events, itemids): if len(measurements_by_time) < 2: return [] - # Convert to sorted list and fill missing indicators - lab_measurements = [] + # Sorted list and fill missing flags + lab_measurements: List[Dict[str, Any]] = [] for days in sorted(measurements_by_time.keys()): - measurement = measurements_by_time[days] - - # Set missing indicators for features not present - measurement.setdefault("missing_egfr", True) - measurement.setdefault("missing_protein", True) - measurement.setdefault("missing_albumin", True) + m = measurements_by_time[days] + m.setdefault("missing_egfr", True) + m.setdefault("missing_protein", True) + m.setdefault("missing_albumin", True) + m.setdefault("egfr", 0.0) + m.setdefault("protein", 0.0) + m.setdefault("albumin", 0.0) + m.setdefault("creatinine", 0.0) + lab_measurements.append(m) - # Set default values for missing features - measurement.setdefault("egfr", 0.0) - measurement.setdefault("protein", 0.0) - measurement.setdefault("albumin", 0.0) - measurement.setdefault("creatinine", 0.0) - - lab_measurements.append(measurement) - - # Collect all missing indicator types present missing_indicators = set() - for measurement in lab_measurements: - for key, value in measurement.items(): + for m in lab_measurements: + for key, value in m.items(): if key.startswith("missing_") and value: missing_indicators.add(key) @@ -408,12 +412,11 @@ def validate_and_convert_lab_events(events, itemids): "lab_measurements": lab_measurements, "missing_indicators": list(missing_indicators), "age": float(age), - "gender": gender, + "gender": [gender], "duration_days": float(duration_days), "has_esrd": has_esrd, } - - return [sample] # Single sample wrapped in list for consistent interface + return [sample] def _calculate_egfr(self, creatinine: float, age: int, gender: str) -> float: """Calculate eGFR using simplified CKD-EPI equation. diff --git a/tests/core/test_mimic4_ckd_surv.py b/tests/core/test_mimic4_ckd_surv.py new file mode 100644 index 000000000..42d90b694 --- /dev/null +++ b/tests/core/test_mimic4_ckd_surv.py @@ -0,0 +1,121 @@ +import os +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path + + +class TestMIMIC4CKDSurv(unittest.TestCase): + """Test CKD survival analysis task on MIMIC-IV demo using ehr_root. + + This test downloads the MIMIC-IV demo dataset from PhysioNet using wget, + constructs a MIMIC4Dataset with ehr_root, and calls set_task() with + MIMIC4CKDSurvAnalysis. The test is tolerant to environments without + network access or wget and will skip gracefully when required resources + or implementations are unavailable. + """ + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.demo_dataset_path = None + self._download_demo_dataset() + self._maybe_import_pyhealth() + + def tearDown(self): + if self.temp_dir and os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def _download_demo_dataset(self): + """Download MIMIC-IV demo dataset using wget (skip if unavailable).""" + download_url = "https://physionet.org/files/mimic-iv-demo/2.2/" + cmd = [ + "wget", + "-r", + "-N", + "-c", + "-np", + "--directory-prefix", + self.temp_dir, + download_url, + ] + + try: + subprocess.run(cmd, check=True, capture_output=True, text=True) + except subprocess.CalledProcessError as e: + raise unittest.SkipTest(f"Failed to download MIMIC-IV demo dataset: {e}") + except FileNotFoundError: + raise unittest.SkipTest( + "wget not available - skipping MIMIC-IV download test" + ) + + # Find the downloaded dataset path + physionet_dir = ( + Path(self.temp_dir) / "physionet.org" / "files" / "mimic-iv-demo" / "2.2" + ) + if physionet_dir.exists(): + self.demo_dataset_path = str(physionet_dir) + else: + raise unittest.SkipTest("Downloaded dataset not found in expected location") + + def _maybe_import_pyhealth(self): + try: + from pyhealth.datasets.mimic4 import MIMIC4Dataset # noqa: F401 + from pyhealth.tasks.ckd_surv import ( # noqa: F401 + MIMIC4CKDSurvAnalysis, + ) + except Exception as e: + raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") + + def test_mimic4_ckd_surv_set_task(self): + """Instantiate MIMIC4Dataset with ehr_root and run set_task().""" + try: + from pyhealth.datasets.mimic4 import MIMIC4Dataset + from pyhealth.tasks.ckd_surv import MIMIC4CKDSurvAnalysis + except Exception as e: + raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") + + # Initialize dataset with EHR tables needed by the task + ehr_tables = [ + "patients", + "admissions", + "labevents", + "diagnoses_icd", + ] + + try: + dataset = MIMIC4Dataset( + ehr_root=self.demo_dataset_path, + ehr_tables=ehr_tables, + dataset_name="mimic4_demo_ehr", + ) + except Exception as e: + raise unittest.SkipTest(f"Failed to construct MIMIC4Dataset: {e}") + + # Build task and run set_task; allow zero samples + # but expect no exceptions during execution + try: + task = MIMIC4CKDSurvAnalysis() + except Exception as e: + raise unittest.SkipTest(f"Task initialization incomplete or failing: {e}") + + try: + sample_dataset = dataset.set_task(task) + except Exception as e: + self.fail(f"set_task() raised an exception: {e}") + + self.assertIsNotNone(sample_dataset, "set_task should return a dataset") + self.assertTrue( + hasattr(sample_dataset, "samples"), + "Returned dataset should have a 'samples' attribute", + ) + + # If samples exist, perform a couple of light checks + if hasattr(sample_dataset, "samples") and sample_dataset.samples: + sample = sample_dataset.samples[0] + self.assertIsInstance(sample, dict) + self.assertIn("patient_id", sample) + + +if __name__ == "__main__": + unittest.main() From be312d7d22009c47d3d97feeeeec1400c7bad581 Mon Sep 17 00:00:00 2001 From: John Wu Date: Fri, 19 Sep 2025 15:16:51 -0500 Subject: [PATCH 4/6] more test cases and logic revisions --- pyhealth/tasks/ckd_surv.py | 482 ++++++++++++++++++----------- tests/core/test_mimic4_ckd_surv.py | 102 ++++-- 2 files changed, 370 insertions(+), 214 deletions(-) diff --git a/pyhealth/tasks/ckd_surv.py b/pyhealth/tasks/ckd_surv.py index 1680977c2..2df857f94 100644 --- a/pyhealth/tasks/ckd_surv.py +++ b/pyhealth/tasks/ckd_surv.py @@ -1,4 +1,5 @@ from typing import Any, Dict, List, Literal, Union, Type +from datetime import timedelta import polars as pl from .base_task import BaseTask @@ -32,7 +33,7 @@ class MIMIC4CKDSurvAnalysis(BaseTask): _FEMALE_GENDER = "F" # Female patients # CKD-EPI 2021 equation constants (from pkgs.data.utils.calculate_eGFR) - _BASE_COEFFICIENT = 141 # Original uses 142 in utils.py + _BASE_COEFFICIENT = 142 # Match original pipeline constant _AGE_FACTOR = 0.993 # Annual age decline factor _FEMALE_ADJUSTMENT = 1.018 # Female gender boost factor @@ -93,13 +94,9 @@ def _configure_schemas( # Use raw processor for time series list of dicts base_input.update({"lab_measurements": RawProcessor}) else: # heterogeneous - # Raw lab measurements with sequence for missing indicators - base_input.update( - { - "lab_measurements": RawProcessor, - "missing_indicators": SequenceProcessor, - } - ) + # Raw lab measurements; per-timestep missing flags inside each + # measurement element + base_input.update({"lab_measurements": RawProcessor}) return base_input, base_output @@ -141,105 +138,234 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if age < self.min_age: return [] - # Get CKD baseline date + # Gather diagnoses ckd_diagnoses = patient.get_events(event_type="diagnoses_icd") ckd_events = [e for e in ckd_diagnoses if e.icd_code in self._CKD_CODES] - if not ckd_events: return [] - - baseline_date = min(e.timestamp for e in ckd_events) - - # Get ESRD outcome esrd_events = [e for e in ckd_diagnoses if e.icd_code in self._ESRD_CODES] + esrd_date = min((e.timestamp for e in esrd_events), default=None) - if esrd_events: - esrd_date = min( - e.timestamp for e in esrd_events if e.timestamp > baseline_date - ) - has_esrd = 1 - duration_days = (esrd_date - baseline_date).days - else: - has_esrd = 0 - # Get all events to find last observation - # FIXED: filter out None timestamps - all_events = patient.get_events() # Get all events - # Filter out events with None timestamps before finding max - valid_events = [e for e in all_events if e.timestamp is not None] - - if valid_events: - last_event = max(valid_events, key=lambda x: x.timestamp) - duration_days = min( - (last_event.timestamp - baseline_date).days, - self.prediction_window_days, + # Collect lab events relevant to the scenario and validate + lab_events = patient.get_events(event_type="labevents") + + def _valid_numeric(e): + try: + return ( + e.valuenum is not None + and float(e.valuenum) > 0 + and e.timestamp is not None + ) + except (ValueError, TypeError): + return False + + # Select labs per scenario + if self.setting in ("time_invariant", "time_variant"): + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS and _valid_numeric(e) + ] + if not creatinine_events: + return [] + # t0 is first creatinine lab + t0 = min(e.timestamp for e in creatinine_events) + # For positives, keep labs up to ESRD date (inclusive by date) + if esrd_date is not None: + # Require at least one lab on the ESRD date to match original + # pipeline + labs_on_esrd_date = [ + e + for e in creatinine_events + if e.timestamp.date() == esrd_date.date() + ] + if not labs_on_esrd_date: + return [] + considered_creatinine = [ + e + for e in creatinine_events + if e.timestamp.date() <= esrd_date.date() + ] + has_esrd = 1 + duration_days = (esrd_date.date() - t0.date()).days + else: + considered_creatinine = creatinine_events + has_esrd = 0 + last_lab_time = max(e.timestamp for e in considered_creatinine) + duration_days = (last_lab_time.date() - t0.date()).days + + # Need at least two labs in the window + if len(considered_creatinine) < 2 or duration_days <= 0: + return [] + + # Dispatch per setting + if self.setting == "time_invariant": + return self._process_time_invariant( + patient, + t0, + age, + gender, + duration_days, + has_esrd, + considered_creatinine, + esrd_date, ) else: - # Fallback: use prediction window if no valid timestamps found - duration_days = self.prediction_window_days + return self._process_time_variant( + patient, + t0, + age, + gender, + duration_days, + has_esrd, + considered_creatinine, + esrd_date, + ) - if duration_days <= 0: - return [] + else: # heterogeneous + # Consider creatinine, protein, albumin + creatinine_events = [ + e + for e in lab_events + if e.itemid in self._CREATININE_ITEMIDS and _valid_numeric(e) + ] + protein_events = [ + e + for e in lab_events + if e.itemid in self._PROTEIN_ITEMIDS and _valid_numeric(e) + ] + albumin_events = [ + e + for e in lab_events + if e.itemid in self._ALBUMIN_ITEMIDS and _valid_numeric(e) + ] + + # Need creatinine to derive egfr at minimum + if not creatinine_events: + return [] + + # t0 is min across all available labs for this scenario + timestamps = [ + e.timestamp + for e in (creatinine_events + protein_events + albumin_events) + if e.timestamp is not None + ] + if not timestamps: + return [] + t0 = min(timestamps) + + if esrd_date is not None: + # Require at least one lab on ESRD date + any_on_esrd = any( + e.timestamp.date() == esrd_date.date() + for e in (creatinine_events + protein_events + albumin_events) + ) + if not any_on_esrd: + return [] + considered_creatinine = [ + e + for e in creatinine_events + if e.timestamp.date() <= esrd_date.date() + ] + considered_protein = [ + e for e in protein_events if e.timestamp.date() <= esrd_date.date() + ] + considered_albumin = [ + e for e in albumin_events if e.timestamp.date() <= esrd_date.date() + ] + has_esrd = 1 + duration_days = (esrd_date.date() - t0.date()).days + else: + considered_creatinine = creatinine_events + considered_protein = protein_events + considered_albumin = albumin_events + has_esrd = 0 + last_time = max( + [ + e.timestamp + for e in ( + considered_creatinine + + considered_protein + + considered_albumin + ) + ] + ) + duration_days = (last_time.date() - t0.date()).days - # Process by setting - if self.setting == "time_invariant": - return self._process_time_invariant( - patient, baseline_date, age, gender, duration_days, has_esrd - ) - elif self.setting == "time_variant": - return self._process_time_variant( - patient, baseline_date, age, gender, duration_days, has_esrd + # Ensure at least two total timepoints across any lab + total_events = len( + { + e.timestamp + for e in ( + considered_creatinine + considered_protein + considered_albumin + ) + } ) - else: # heterogeneous + if total_events < 2 or duration_days <= 0: + return [] + return self._process_heterogeneous( - patient, baseline_date, age, gender, duration_days, has_esrd + patient, + t0, + age, + gender, + duration_days, + has_esrd, + considered_creatinine, + considered_protein, + considered_albumin, + esrd_date, ) def _process_time_invariant( - self, patient, baseline_date, age, gender, duration_days, has_esrd + self, + patient, + t0, + age, + gender, + duration_days, + has_esrd, + considered_creatinine, + esrd_date, ): - """Process for time-invariant analysis.""" - lab_events = patient.get_events(event_type="labevents") - creatinine_events = [ - e - for e in lab_events - if ( - e.itemid in self._CREATININE_ITEMIDS - and e.valuenum is not None - and e.timestamp >= baseline_date - ) - ] - - if not creatinine_events: - return [] + """ + Process for time-invariant analysis aligned with original + NON_TIME_VARIANT. - # Validate and find baseline creatinine - valid_creatinine_events = [] - for e in creatinine_events: - try: - creatinine_value = float(e.valuenum) - if creatinine_value > 0: - valid_creatinine_events.append((e, creatinine_value)) - except (ValueError, TypeError): - continue + - Positives: pick lab on ESRD date (last that day) and compute egfr + - Negatives: pick last available lab + """ + # Choose target creatinine event + if has_esrd and esrd_date is not None: + same_day_events = [ + e + for e in considered_creatinine + if e.timestamp.date() == esrd_date.date() + ] + if not same_day_events: + return [] + target_event = max(same_day_events, key=lambda x: x.timestamp) + else: + target_event = max(considered_creatinine, key=lambda x: x.timestamp) - if not valid_creatinine_events: + try: + creatinine_value = float(target_event.valuenum) + except (ValueError, TypeError): + return [] + if creatinine_value <= 0: return [] - # Closest to baseline - _, baseline_creatinine_value = min( - valid_creatinine_events, - key=lambda x: abs((x[0].timestamp - baseline_date).days), - ) - - egfr = self._calculate_egfr(baseline_creatinine_value, age, gender) + egfr = self._calculate_egfr(creatinine_value, age, gender) - # Comorbidities before baseline + # Comorbidities before first lab (t0) diagnoses = patient.get_events(event_type="diagnoses_icd") comorbidities = [ - e.icd_code for e in diagnoses if e.timestamp <= baseline_date and e.icd_code + e.icd_code + for e in diagnoses + if e.timestamp is not None and e.timestamp <= t0 and e.icd_code ] - # Race from admissions + # Race from admissions (optional meta) admissions = patient.get_events(event_type="admissions") race = admissions[0].race if admissions else "unknown" @@ -259,28 +385,26 @@ def _process_time_invariant( return [sample] def _process_time_variant( - self, patient, baseline_date, age, gender, duration_days, has_esrd + self, + patient, + t0, + age, + gender, + duration_days, + has_esrd, + considered_creatinine, + esrd_date, ): - """Process for time-varying analysis.""" - lab_events = patient.get_events(event_type="labevents") - creatinine_events = [ - e - for e in lab_events - if ( - e.itemid in self._CREATININE_ITEMIDS - and e.valuenum is not None - and e.timestamp >= baseline_date - ) - ] - - if len(creatinine_events) < 2: - return [] + """ + Process for time-varying analysis aligned with original + TIME_VARIANT. - # Chronological order - creatinine_events.sort(key=lambda x: x.timestamp) + Build series from first lab (t0) up to ESRD date (if positive) or last + lab (negative). + """ + considered_creatinine.sort(key=lambda x: x.timestamp) lab_measurements = [] - - for e in creatinine_events: + for e in considered_creatinine: try: creatinine_value = float(e.valuenum) if creatinine_value <= 0: @@ -288,15 +412,16 @@ def _process_time_variant( except (ValueError, TypeError): continue - days_from_baseline = (e.timestamp - baseline_date).days + days_from_t0 = (e.timestamp.date() - t0.date()).days egfr_value = self._calculate_egfr(creatinine_value, age, gender) - lab_measurements.append( - { - "timestamp": days_from_baseline, - "egfr": egfr_value, - "creatinine": creatinine_value, - } - ) + m = { + "timestamp": days_from_t0, + "egfr": egfr_value, + "creatinine": creatinine_value, + } + if esrd_date is not None: + m["has_esrd_step"] = int(e.timestamp.date() == esrd_date.date()) + lab_measurements.append(m) age_group = "elderly" if age >= 65 else "adult" gender_str = "male" if gender == self._MALE_GENDER else "female" @@ -313,96 +438,84 @@ def _process_time_variant( return [sample] def _process_heterogeneous( - self, patient, baseline_date, age, gender, duration_days, has_esrd + self, + patient, + t0, + age, + gender, + duration_days, + has_esrd, + creatinine_events, + protein_events, + albumin_events, + esrd_date, ): - """Process for heterogeneous analysis with missing indicators.""" - lab_events = patient.get_events(event_type="labevents") - - # Validator for lab values - def validate(events, itemids): - out = [] - for e in events: - if e.itemid in itemids and e.valuenum is not None: - try: - v = float(e.valuenum) - if v > 0: - out.append((e, v)) - except (ValueError, TypeError): - continue - return out - - creatinine_events = validate(lab_events, self._CREATININE_ITEMIDS) - protein_events = validate(lab_events, self._PROTEIN_ITEMIDS) - albumin_events = validate(lab_events, self._ALBUMIN_ITEMIDS) - - if not creatinine_events: - return [] + """Process for heterogeneous analysis with per-timestep missing flags. + Missing flags use names: egfr_missing, protein_missing, albumin_missing + (0/1). + """ measurements_by_time: Dict[int, Dict[str, Any]] = {} - # Creatinine/eGFR - for e, creatinine_value in creatinine_events: - if e.timestamp >= baseline_date: - days = (e.timestamp - baseline_date).days - egfr = self._calculate_egfr(creatinine_value, age, gender) - if days not in measurements_by_time: - measurements_by_time[days] = {"timestamp": days} - measurements_by_time[days].update( - { - "egfr": egfr, - "creatinine": creatinine_value, - "missing_egfr": False, - } - ) + def _upsert(days: int, updates: Dict[str, Any]): + if days not in measurements_by_time: + measurements_by_time[days] = { + "timestamp": days, + "egfr_missing": 1, + "protein_missing": 1, + "albumin_missing": 1, + "egfr": 0.0, + "protein": 0.0, + "albumin": 0.0, + "creatinine": 0.0, + } + measurements_by_time[days].update(updates) - # Protein - for e, protein_value in protein_events: - if e.timestamp >= baseline_date: - days = (e.timestamp - baseline_date).days - if days not in measurements_by_time: - measurements_by_time[days] = { - "timestamp": days, - "missing_egfr": True, - } - measurements_by_time[days].update( - {"protein": protein_value, "missing_protein": False} - ) + # Within-window events already considered by caller + for e in creatinine_events: + days = (e.timestamp.date() - t0.date()).days + try: + cr = float(e.valuenum) + except (ValueError, TypeError): + continue + if cr <= 0: + continue + egfr = self._calculate_egfr(cr, age, gender) + _upsert(days, {"egfr": egfr, "creatinine": cr, "egfr_missing": 0}) - # Albumin - for e, albumin_value in albumin_events: - if e.timestamp >= baseline_date: - days = (e.timestamp - baseline_date).days - if days not in measurements_by_time: - measurements_by_time[days] = { - "timestamp": days, - "missing_egfr": True, - } - measurements_by_time[days].update( - {"albumin": albumin_value, "missing_albumin": False} - ) + for e in protein_events: + days = (e.timestamp.date() - t0.date()).days + try: + pv = float(e.valuenum) + except (ValueError, TypeError): + continue + if pv <= 0: + continue + _upsert(days, {"protein": pv, "protein_missing": 0}) + + for e in albumin_events: + days = (e.timestamp.date() - t0.date()).days + try: + av = float(e.valuenum) + except (ValueError, TypeError): + continue + if av <= 0: + continue + _upsert(days, {"albumin": av, "albumin_missing": 0}) if len(measurements_by_time) < 2: return [] - # Sorted list and fill missing flags lab_measurements: List[Dict[str, Any]] = [] for days in sorted(measurements_by_time.keys()): m = measurements_by_time[days] - m.setdefault("missing_egfr", True) - m.setdefault("missing_protein", True) - m.setdefault("missing_albumin", True) - m.setdefault("egfr", 0.0) - m.setdefault("protein", 0.0) - m.setdefault("albumin", 0.0) - m.setdefault("creatinine", 0.0) + if esrd_date is not None: + # Set step-level ESRD flag when day matches ESRD date + m["has_esrd_step"] = int( + (t0.date() + timedelta(days=days)) == esrd_date.date() + ) lab_measurements.append(m) - missing_indicators = set() - for m in lab_measurements: - for key, value in m.items(): - if key.startswith("missing_") and value: - missing_indicators.add(key) - age_group = "elderly" if age >= 65 else "adult" gender_str = "male" if gender == self._MALE_GENDER else "female" @@ -410,7 +523,6 @@ def validate(events, itemids): "patient_id": patient.patient_id, "demographics": [age_group, gender_str], "lab_measurements": lab_measurements, - "missing_indicators": list(missing_indicators), "age": float(age), "gender": [gender], "duration_days": float(duration_days), diff --git a/tests/core/test_mimic4_ckd_surv.py b/tests/core/test_mimic4_ckd_surv.py index 42d90b694..e3a3d9bf9 100644 --- a/tests/core/test_mimic4_ckd_surv.py +++ b/tests/core/test_mimic4_ckd_surv.py @@ -4,6 +4,7 @@ import tempfile import unittest from pathlib import Path +import json class TestMIMIC4CKDSurv(unittest.TestCase): @@ -43,7 +44,9 @@ def _download_demo_dataset(self): try: subprocess.run(cmd, check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: - raise unittest.SkipTest(f"Failed to download MIMIC-IV demo dataset: {e}") + raise unittest.SkipTest( + f"Failed to download MIMIC-IV demo dataset: {e}" + ) except FileNotFoundError: raise unittest.SkipTest( "wget not available - skipping MIMIC-IV download test" @@ -51,12 +54,18 @@ def _download_demo_dataset(self): # Find the downloaded dataset path physionet_dir = ( - Path(self.temp_dir) / "physionet.org" / "files" / "mimic-iv-demo" / "2.2" + Path(self.temp_dir) + / "physionet.org" + / "files" + / "mimic-iv-demo" + / "2.2" ) if physionet_dir.exists(): self.demo_dataset_path = str(physionet_dir) else: - raise unittest.SkipTest("Downloaded dataset not found in expected location") + raise unittest.SkipTest( + "Downloaded dataset not found in expected location" + ) def _maybe_import_pyhealth(self): try: @@ -65,15 +74,24 @@ def _maybe_import_pyhealth(self): MIMIC4CKDSurvAnalysis, ) except Exception as e: - raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") + raise unittest.SkipTest( + f"pyhealth import failed or incomplete: {e}" + ) + + def test_mimic4_ckd_surv_all_modes(self): + """Instantiate MIMIC4Dataset and run set_task() for all modes. - def test_mimic4_ckd_surv_set_task(self): - """Instantiate MIMIC4Dataset with ehr_root and run set_task().""" + Modes: time_invariant, time_variant, heterogeneous. For each mode, + assert dataset is constructed and, if samples exist, print the first + sample for visualization. + """ try: from pyhealth.datasets.mimic4 import MIMIC4Dataset from pyhealth.tasks.ckd_surv import MIMIC4CKDSurvAnalysis except Exception as e: - raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") + raise unittest.SkipTest( + f"pyhealth import failed or incomplete: {e}" + ) # Initialize dataset with EHR tables needed by the task ehr_tables = [ @@ -92,29 +110,55 @@ def test_mimic4_ckd_surv_set_task(self): except Exception as e: raise unittest.SkipTest(f"Failed to construct MIMIC4Dataset: {e}") - # Build task and run set_task; allow zero samples - # but expect no exceptions during execution - try: - task = MIMIC4CKDSurvAnalysis() - except Exception as e: - raise unittest.SkipTest(f"Task initialization incomplete or failing: {e}") - - try: - sample_dataset = dataset.set_task(task) - except Exception as e: - self.fail(f"set_task() raised an exception: {e}") - - self.assertIsNotNone(sample_dataset, "set_task should return a dataset") - self.assertTrue( - hasattr(sample_dataset, "samples"), - "Returned dataset should have a 'samples' attribute", - ) + modes = ["time_invariant", "time_variant", "heterogeneous"] + for mode in modes: + try: + task = MIMIC4CKDSurvAnalysis(setting=mode) + except Exception as e: + raise unittest.SkipTest( + f"Task initialization failed for mode '{mode}': {e}" + ) + + try: + sample_dataset = dataset.set_task(task) + except Exception as e: + self.fail( + f"set_task() raised an exception in mode '{mode}': {e}" + ) + + self.assertIsNotNone( + sample_dataset, "set_task should return a dataset" + ) + self.assertTrue( + hasattr(sample_dataset, "samples"), + "Returned dataset should have a 'samples' attribute", + ) - # If samples exist, perform a couple of light checks - if hasattr(sample_dataset, "samples") and sample_dataset.samples: - sample = sample_dataset.samples[0] - self.assertIsInstance(sample, dict) - self.assertIn("patient_id", sample) + n = len(getattr(sample_dataset, "samples", [])) + print(f"\n[ckd_surv] mode={mode} samples={n}") + + if n: + sample = sample_dataset.samples[0] + self.assertIsInstance(sample, dict) + self.assertIn("patient_id", sample) + + # Pretty-print a compact view depending on mode + view = {k: sample.get(k) for k in ( + "patient_id", + "duration_days", + "has_esrd", + "baseline_egfr", + "lab_measurements", + ) if k in sample} + + # Truncate lab_measurements if long + if "lab_measurements" in view and isinstance( + view["lab_measurements"], list + ): + lm = view["lab_measurements"] + view["lab_measurements"] = lm[:3] + + print(json.dumps(view, indent=2, default=str)) if __name__ == "__main__": From 0f79bf1a8a2d9615bf4ee2e480192039d9f8a870 Mon Sep 17 00:00:00 2001 From: John Wu Date: Fri, 19 Sep 2025 15:17:09 -0500 Subject: [PATCH 5/6] test cases --- tests/core/test_mimic4_ckd_surv.py | 48 +++++++++++------------------- 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/tests/core/test_mimic4_ckd_surv.py b/tests/core/test_mimic4_ckd_surv.py index e3a3d9bf9..34b3b8331 100644 --- a/tests/core/test_mimic4_ckd_surv.py +++ b/tests/core/test_mimic4_ckd_surv.py @@ -44,9 +44,7 @@ def _download_demo_dataset(self): try: subprocess.run(cmd, check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: - raise unittest.SkipTest( - f"Failed to download MIMIC-IV demo dataset: {e}" - ) + raise unittest.SkipTest(f"Failed to download MIMIC-IV demo dataset: {e}") except FileNotFoundError: raise unittest.SkipTest( "wget not available - skipping MIMIC-IV download test" @@ -54,18 +52,12 @@ def _download_demo_dataset(self): # Find the downloaded dataset path physionet_dir = ( - Path(self.temp_dir) - / "physionet.org" - / "files" - / "mimic-iv-demo" - / "2.2" + Path(self.temp_dir) / "physionet.org" / "files" / "mimic-iv-demo" / "2.2" ) if physionet_dir.exists(): self.demo_dataset_path = str(physionet_dir) else: - raise unittest.SkipTest( - "Downloaded dataset not found in expected location" - ) + raise unittest.SkipTest("Downloaded dataset not found in expected location") def _maybe_import_pyhealth(self): try: @@ -74,9 +66,7 @@ def _maybe_import_pyhealth(self): MIMIC4CKDSurvAnalysis, ) except Exception as e: - raise unittest.SkipTest( - f"pyhealth import failed or incomplete: {e}" - ) + raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") def test_mimic4_ckd_surv_all_modes(self): """Instantiate MIMIC4Dataset and run set_task() for all modes. @@ -89,9 +79,7 @@ def test_mimic4_ckd_surv_all_modes(self): from pyhealth.datasets.mimic4 import MIMIC4Dataset from pyhealth.tasks.ckd_surv import MIMIC4CKDSurvAnalysis except Exception as e: - raise unittest.SkipTest( - f"pyhealth import failed or incomplete: {e}" - ) + raise unittest.SkipTest(f"pyhealth import failed or incomplete: {e}") # Initialize dataset with EHR tables needed by the task ehr_tables = [ @@ -122,13 +110,9 @@ def test_mimic4_ckd_surv_all_modes(self): try: sample_dataset = dataset.set_task(task) except Exception as e: - self.fail( - f"set_task() raised an exception in mode '{mode}': {e}" - ) + self.fail(f"set_task() raised an exception in mode '{mode}': {e}") - self.assertIsNotNone( - sample_dataset, "set_task should return a dataset" - ) + self.assertIsNotNone(sample_dataset, "set_task should return a dataset") self.assertTrue( hasattr(sample_dataset, "samples"), "Returned dataset should have a 'samples' attribute", @@ -143,13 +127,17 @@ def test_mimic4_ckd_surv_all_modes(self): self.assertIn("patient_id", sample) # Pretty-print a compact view depending on mode - view = {k: sample.get(k) for k in ( - "patient_id", - "duration_days", - "has_esrd", - "baseline_egfr", - "lab_measurements", - ) if k in sample} + view = { + k: sample.get(k) + for k in ( + "patient_id", + "duration_days", + "has_esrd", + "baseline_egfr", + "lab_measurements", + ) + if k in sample + } # Truncate lab_measurements if long if "lab_measurements" in view and isinstance( From 130d60b12eb514a7faa31e9f3196798079b48378 Mon Sep 17 00:00:00 2001 From: John Wu Date: Fri, 19 Sep 2025 15:46:53 -0500 Subject: [PATCH 6/6] documentation updates --- pyhealth/tasks/ckd_surv.py | 208 ++++++++++++++++++++++++++++++++++--- 1 file changed, 191 insertions(+), 17 deletions(-) diff --git a/pyhealth/tasks/ckd_surv.py b/pyhealth/tasks/ckd_surv.py index 2df857f94..deaf7c985 100644 --- a/pyhealth/tasks/ckd_surv.py +++ b/pyhealth/tasks/ckd_surv.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Literal, Union, Type +from typing import Any, Dict, List, Literal, Union, Type, Optional from datetime import timedelta import polars as pl @@ -11,13 +11,100 @@ class MIMIC4CKDSurvAnalysis(BaseTask): - """CKD survival analysis task with simplified configuration. + """Survival analysis for CKD progression on MIMIC-IV (CKD -> ESRD). + + This task prepares patient-level samples for survival modeling using + MIMIC-IV tables (patients, admissions, diagnoses_icd, labevents). It + supports three settings that change the input form: + + - "time_invariant": single-row snapshot per patient + - "time_variant": time series with a single modality stream + - "heterogeneous": time series with multiple lab modalities + + The time origin (t0) is the first available lab in the window. + Positive cases are censored at the ESRD date (inclusive-by-date) and + negatives at the last available lab. Durations are computed in days + from t0. + + Inputs and outputs by setting + - Common output (all settings): + - duration_days: Tensor (float), days between t0 and censoring + - has_esrd: Tensor (int, 0/1), whether ESRD occurred in the window + + - time_invariant inputs: + - demographics: Sequence ([age_group, gender_str, race]) + - age: Tensor (float) + - gender: Sequence (["M"|"F"]) for modeling as categorical + - baseline_egfr: Tensor (float), from a single target lab + - comorbidities: Sequence (ICD codes prior to t0) + + - time_variant inputs: + - demographics: Sequence ([age_group, gender_str]) + - age: Tensor (float) + - gender: Sequence (["M"|"F"]) for modeling as categorical + - lab_measurements: Raw list[dict], ordered by days since t0 + Each element includes: + - timestamp: int, days since t0 + - creatinine: float (if present) + - egfr: float (if present) + - has_esrd_step: int (0/1), only when ESRD-day exists + - extras via extra_lab_itemids (e.g. bun) + - bun_missing flag (0 present, 1 missing) + + - heterogeneous inputs: + - demographics, age, gender: same as time_variant + - lab_measurements: Raw list[dict] with multimodal labs per day + Each element includes (when present): + - timestamp: int + - creatinine: float + - egfr: float, derived from creatinine, age, gender + - protein: float, albumin: float + - egfr_missing/protein_missing/albumin_missing: int (0/1) + - has_esrd_step: int (0/1) on the ESRD day + - Any configured extras plus {name}_missing flags + + Parameters + - setting: one of ["time_invariant", "time_variant", "heterogeneous"] + - min_age: minimum age (years) to include in cohort (default 18) + - prediction_window_days: not used to truncate currently; reserved + - extra_lab_itemids: optional dict mapping feature name -> list of + labevents.itemid strings to include as extra modalities. For each + name, + two fields may appear in lab_measurements: {name} (float) and + {name}_missing (int 0/1). Values are aligned to days since t0. + + Notes + - eGFR uses the CKD-EPI 2021 formula with base coefficient 142. See + https://pubmed.ncbi.nlm.nih.gov/34554658/ + - Positives require at least one lab event recorded on the ESRD date, + matching the original pipeline semantics. + + Example + ------- + >>> from pyhealth.datasets import MIMIC4Dataset + >>> from pyhealth.tasks.ckd_surv import MIMIC4CKDSurvAnalysis + >>> dataset = MIMIC4Dataset( + ... root="/path/to/mimiciv/demo", + ... tables=[ + ... "patients", "admissions", "labevents", "diagnoses_icd" + ... ], + ... dev=True, + ... ) + >>> task = MIMIC4CKDSurvAnalysis( + ... setting="time_variant", + ... extra_lab_itemids={"bun": ["51006"]}, + ... ) + >>> dataset.set_task(task) + >>> samples = dataset.samples + >>> sample = samples[0] + >>> sorted(sample.keys()) + ['age', 'demographics', 'duration_days', 'gender', 'has_esrd', + 'lab_measurements', 'patient_id'] + >>> sample['lab_measurements'][0].keys() + dict_keys(['timestamp', 'egfr_missing', 'protein_missing', + 'albumin_missing', 'egfr', 'protein', 'albumin', + 'creatinine', 'has_esrd_step', 'bun_missing', 'bun']) - eGFR calculation methodology adapted from: - - Original implementation: pkgs.data.utils.calculate_eGFR() - - Formula source: pkgs.data.store.get_egfr_df() - - Reference: CKD-EPI 2021 formula - (https://pubmed.ncbi.nlm.nih.gov/34554658/) """ # Private class variables for settings @@ -51,6 +138,7 @@ def __init__( ] = "time_invariant", min_age: int = 18, prediction_window_days: int = 365 * 5, + extra_lab_itemids: Optional[Dict[str, List[str]]] = None, ): if setting not in self._SURVIVAL_SETTINGS: @@ -60,6 +148,11 @@ def __init__( self.min_age = min_age self.prediction_window_days = prediction_window_days self.task_name = f"MIMIC4CKDSurvAnalysis_{self.setting}" + # Optional extensibility: additional lab item IDs to extract from + # labevents. Dict maps feature name -> list of itemids. Values will + # appear inside lab_measurements as the feature name and a + # corresponding "{feature}_missing" flag (0 present, 1 missing). + self.extra_lab_itemids = extra_lab_itemids or {} self.input_schema, self.output_schema = self._configure_schemas() def _configure_schemas( @@ -159,6 +252,15 @@ def _valid_numeric(e): except (ValueError, TypeError): return False + # Pre-compute extra lab events (validated numeric) + extra_events_map: Dict[str, List[Any]] = {} + for feat_name, itemids in self.extra_lab_itemids.items(): + if not itemids: + continue + extra_events_map[feat_name] = [ + e for e in lab_events if e.itemid in itemids and _valid_numeric(e) + ] + # Select labs per scenario if self.setting in ("time_invariant", "time_variant"): creatinine_events = [ @@ -211,6 +313,16 @@ def _valid_numeric(e): esrd_date, ) else: + # Filter extras by ESRD date if positive + if has_esrd and esrd_date is not None: + filtered_extras: Dict[str, List[Any]] = {} + for name, events in extra_events_map.items(): + filtered_extras[name] = [ + e for e in events if e.timestamp.date() <= esrd_date.date() + ] + else: + filtered_extras = extra_events_map + return self._process_time_variant( patient, t0, @@ -220,6 +332,7 @@ def _valid_numeric(e): has_esrd, considered_creatinine, esrd_date, + filtered_extras, ) else: # heterogeneous @@ -247,7 +360,12 @@ def _valid_numeric(e): # t0 is min across all available labs for this scenario timestamps = [ e.timestamp - for e in (creatinine_events + protein_events + albumin_events) + for e in ( + creatinine_events + + protein_events + + albumin_events + + [ev for lst in extra_events_map.values() for ev in lst] + ) if e.timestamp is not None ] if not timestamps: @@ -258,7 +376,12 @@ def _valid_numeric(e): # Require at least one lab on ESRD date any_on_esrd = any( e.timestamp.date() == esrd_date.date() - for e in (creatinine_events + protein_events + albumin_events) + for e in ( + creatinine_events + + protein_events + + albumin_events + + [ev for lst in extra_events_map.values() for ev in lst] + ) ) if not any_on_esrd: return [] @@ -273,12 +396,17 @@ def _valid_numeric(e): considered_albumin = [ e for e in albumin_events if e.timestamp.date() <= esrd_date.date() ] + considered_extras = { + name: [e for e in events if e.timestamp.date() <= esrd_date.date()] + for name, events in extra_events_map.items() + } has_esrd = 1 duration_days = (esrd_date.date() - t0.date()).days else: considered_creatinine = creatinine_events considered_protein = protein_events considered_albumin = albumin_events + considered_extras = extra_events_map has_esrd = 0 last_time = max( [ @@ -287,6 +415,7 @@ def _valid_numeric(e): considered_creatinine + considered_protein + considered_albumin + + [ev for lst in considered_extras.values() for ev in lst] ) ] ) @@ -297,7 +426,10 @@ def _valid_numeric(e): { e.timestamp for e in ( - considered_creatinine + considered_protein + considered_albumin + considered_creatinine + + considered_protein + + considered_albumin + + [ev for lst in considered_extras.values() for ev in lst] ) } ) @@ -315,6 +447,7 @@ def _valid_numeric(e): considered_protein, considered_albumin, esrd_date, + considered_extras, ) def _process_time_invariant( @@ -394,6 +527,7 @@ def _process_time_variant( has_esrd, considered_creatinine, esrd_date, + extra_events_map: Optional[Dict[str, List[Any]]] = None, ): """ Process for time-varying analysis aligned with original @@ -402,8 +536,19 @@ def _process_time_variant( Build series from first lab (t0) up to ESRD date (if positive) or last lab (negative). """ + # Build union of timepoints across creatinine and extras + extra_events_map = extra_events_map or {} + measurements_by_time: Dict[int, Dict[str, Any]] = {} + + def _ensure_day(day: int): + if day not in measurements_by_time: + measurements_by_time[day] = {"timestamp": day} + for name in extra_events_map.keys(): + measurements_by_time[day][f"{name}_missing"] = 1 + measurements_by_time[day][name] = 0.0 + + # Creatinine and egfr considered_creatinine.sort(key=lambda x: x.timestamp) - lab_measurements = [] for e in considered_creatinine: try: creatinine_value = float(e.valuenum) @@ -414,14 +559,29 @@ def _process_time_variant( days_from_t0 = (e.timestamp.date() - t0.date()).days egfr_value = self._calculate_egfr(creatinine_value, age, gender) - m = { - "timestamp": days_from_t0, - "egfr": egfr_value, - "creatinine": creatinine_value, - } + _ensure_day(days_from_t0) + m = measurements_by_time[days_from_t0] + m["egfr"] = egfr_value + m["creatinine"] = creatinine_value if esrd_date is not None: m["has_esrd_step"] = int(e.timestamp.date() == esrd_date.date()) - lab_measurements.append(m) + # Extras + for name, events in extra_events_map.items(): + for e in events: + day = (e.timestamp.date() - t0.date()).days + try: + val = float(e.valuenum) + except (ValueError, TypeError): + continue + if val <= 0: + continue + _ensure_day(day) + measurements_by_time[day][name] = val + measurements_by_time[day][f"{name}_missing"] = 0 + + lab_measurements = [ + measurements_by_time[d] for d in sorted(measurements_by_time.keys()) + ] age_group = "elderly" if age >= 65 else "adult" gender_str = "male" if gender == self._MALE_GENDER else "female" @@ -449,6 +609,7 @@ def _process_heterogeneous( protein_events, albumin_events, esrd_date, + extra_events_map: Optional[Dict[str, List[Any]]] = None, ): """Process for heterogeneous analysis with per-timestep missing flags. @@ -456,6 +617,7 @@ def _process_heterogeneous( (0/1). """ measurements_by_time: Dict[int, Dict[str, Any]] = {} + extra_events_map = extra_events_map or {} def _upsert(days: int, updates: Dict[str, Any]): if days not in measurements_by_time: @@ -503,6 +665,18 @@ def _upsert(days: int, updates: Dict[str, Any]): continue _upsert(days, {"albumin": av, "albumin_missing": 0}) + # Extras + for name, events in extra_events_map.items(): + for e in events: + days = (e.timestamp.date() - t0.date()).days + try: + val = float(e.valuenum) + except (ValueError, TypeError): + continue + if val <= 0: + continue + _upsert(days, {name: val, f"{name}_missing": 0}) + if len(measurements_by_time) < 2: return []