diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index bdaa9599a..e2a8b8a02 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -207,6 +207,7 @@ Available Tasks Base Task In-Hospital Mortality (MIMIC-IV) + Sepsis Prediction (MIMIC-IV) In-Hospital Mortality (MEDS) MIMIC-III ICD-9 Coding Cardiology Detection diff --git a/docs/api/tasks/pyhealth.tasks.SepsisPredictionMIMIC4.rst b/docs/api/tasks/pyhealth.tasks.SepsisPredictionMIMIC4.rst new file mode 100644 index 000000000..1720186ab --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.SepsisPredictionMIMIC4.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.SepsisPredictionMIMIC4 +======================================= + +.. autoclass:: pyhealth.tasks.sepsis_prediction_mimic4.SepsisPredictionMIMIC4 + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/sepsis_prediction_mimic4_demo.py b/examples/sepsis_prediction_mimic4_demo.py new file mode 100644 index 000000000..174397daf --- /dev/null +++ b/examples/sepsis_prediction_mimic4_demo.py @@ -0,0 +1,76 @@ +# Author: Anish Gupta +# NetID: anishg8 +# Paper Title: N/A (original task contribution, not a paper reproduction) +# Paper Link: N/A +# Description: End-to-end example running SepsisPredictionMIMIC4 on real +# MIMIC-IV data: load the dataset, build the task's sample dataset, +# split by patient, and train/evaluate a small RNN. +"""End-to-end example: sepsis prediction on MIMIC-IV with PyHealth. + +Requires access to MIMIC-IV (PhysioNet credentialing: +https://physionet.org/content/mimiciv/), including its ICU module +(``icu/chartevents.csv.gz``, ``icu/d_items.csv.gz``) for vitals. + +Run: + + python examples/sepsis_prediction_mimic4_demo.py --root /path/to/mimic-iv/2.2 + +See ``pyhealth.tasks.SepsisPredictionMIMIC4`` for the label definition +(qSOFA-based Sepsis-3 approximation) and its documented limitations. +""" + +import argparse + +from pyhealth.datasets import MIMIC4EHRDataset, get_dataloader, split_by_patient +from pyhealth.models import RNN +from pyhealth.tasks import SepsisPredictionMIMIC4 +from pyhealth.trainer import Trainer + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + required=True, + help="Root of the MIMIC-IV dataset (directory containing hosp/ and icu/)", + ) + args = parser.parse_args() + + dataset = MIMIC4EHRDataset( + root=args.root, + tables=["admissions", "prescriptions", "labevents", "chartevents"], + ) + dataset.stats() + + sample_dataset = dataset.set_task(SepsisPredictionMIMIC4()) + n_positive = sum(int(s["sepsis"]) for s in sample_dataset) + print( + f"Sepsis samples: {len(sample_dataset)} total, {n_positive} positive " + f"({n_positive / max(len(sample_dataset), 1):.1%})" + ) + + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True) + val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False) + test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False) + + model = RNN(dataset=sample_dataset, hidden_dim=64) + trainer = Trainer(model=model) + trainer.train( + train_dataloader=train_dataloader, + val_dataloader=val_dataloader, + epochs=5, + monitor="pr_auc", + ) + # Sepsis is a rare-outcome task: report pr_auc/roc_auc, not just + # accuracy, which would be misleadingly high for a model that mostly + # predicts the majority class. + metrics = trainer.evaluate(test_dataloader) + print("Test metrics:", metrics) + + +if __name__ == "__main__": + # BaseDataset spawns Dask worker processes; keep the main-module guard. + main() diff --git a/pyhealth/datasets/configs/mimic4_ehr.yaml b/pyhealth/datasets/configs/mimic4_ehr.yaml index 84c570bb9..4f6def983 100644 --- a/pyhealth/datasets/configs/mimic4_ehr.yaml +++ b/pyhealth/datasets/configs/mimic4_ehr.yaml @@ -117,3 +117,25 @@ tables: - "hcpcs_cd" - "seq_num" - "short_description" + + chartevents: + file_path: "icu/chartevents.csv.gz" + patient_id: "subject_id" + join: + - file_path: "icu/d_items.csv.gz" + "on": "itemid" + how: "inner" + columns: + - "label" + - "category" + timestamp: "charttime" + attributes: + - "hadm_id" + - "stay_id" + - "itemid" + - "label" + - "category" + - "value" + - "valuenum" + - "valueuom" + - "storetime" diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index df8411db0..4648da46a 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -56,6 +56,9 @@ to_evaluation_dataframe, ) from .patient_linkage import patient_linkage_mimic3_fn +from .sepsis_prediction_mimic4 import ( + SepsisPredictionMIMIC4 as SepsisPredictionMIMIC4, +) from .readmission_prediction import ( ReadmissionPredictionEICU, ReadmissionPredictionMIMIC3, diff --git a/pyhealth/tasks/sepsis_prediction_mimic4.py b/pyhealth/tasks/sepsis_prediction_mimic4.py new file mode 100644 index 000000000..a2134557c --- /dev/null +++ b/pyhealth/tasks/sepsis_prediction_mimic4.py @@ -0,0 +1,368 @@ +# Author: Anish Gupta +# NetID: anishg8 +# Paper Title: N/A (original task contribution, not a paper reproduction) +# Paper Link: N/A +# Description: Sepsis prediction task for MIMIC-IV. Labels each hospital +# admission for sepsis onset using a qSOFA-based approximation of the +# Sepsis-3 clinical criteria (Seymour et al., "Assessment of Clinical +# Criteria for Sepsis," JAMA 2016, https://doi.org/10.1001/jama.2016.0288), +# built on the labs/vitals/prescriptions tables PyHealth's MIMIC-IV +# loader already supports (plus a new `chartevents` table added to +# mimic4_ehr.yaml for vitals). + +from datetime import datetime, timedelta +from typing import Any, ClassVar + +import polars as pl + +from .base_task import BaseTask + + +class SepsisPredictionMIMIC4(BaseTask): + """Task for predicting sepsis onset during a MIMIC-IV hospital admission. + + Each admission is labeled using a Sepsis-3-style two-signal definition: + a *suspected infection* signal (a new antibiotic order) co-occurring + with an *organ-dysfunction* signal (qSOFA >= 2). + + qSOFA ("quick SOFA") is the bedside, vitals-only simplification of the + full Sepsis-3 SOFA score, from the same consensus definitions (Seymour + et al., "Assessment of Clinical Criteria for Sepsis," JAMA 2016). It + flags a patient when at least 2 of the following 3 criteria are met: + respiratory rate >= 22/min, systolic blood pressure <= 100 mmHg, and + altered mental status (Glasgow Coma Scale < 15). Full SOFA additionally + requires vasopressor dosing and urine output, which are not currently + loaded by PyHealth's MIMIC-IV configuration, so qSOFA is used here as a + documented, citable simplification rather than full Sepsis-3. + + Limitations: + - qSOFA is a simplified organ-dysfunction proxy, not the full SOFA + score; it will disagree with a full Sepsis-3 label in some cases. + - "Suspected infection" is approximated by a curated antibiotic + drug-name whitelist (``ANTIBIOTIC_DRUG_NAMES``) matched against + the free-text ``prescriptions.drug`` field. This is a heuristic: + it will miss antibiotics not on the list and may match ambiguous + names. + - Onset is the earliest timestamp where both signals co-occur + within ``ANTIBIOTIC_WINDOW_HOURS`` of each other. Features are + censored strictly before that timestamp so the triggering + observation itself is never part of the model's input. + - The ``chartevents``/``labevents`` item IDs below are the standard + MIMIC-IV IDs for these measurements but should be re-verified + against a real ``icu/d_items.csv.gz`` / ``hosp/d_labitems.csv.gz`` + before use on production data. + + Attributes: + task_name: The name of the task. + input_schema: ``{"observations": "timeseries"}`` -- a combined + labs + vitals timeseries, censored before any sepsis onset. + output_schema: ``{"sepsis": "binary"}``. + + Examples: + >>> from pyhealth.datasets import MIMIC4EHRDataset + >>> from pyhealth.tasks import SepsisPredictionMIMIC4 + >>> dataset = MIMIC4EHRDataset( + ... root="/path/to/mimic-iv/2.2", + ... tables=["admissions", "prescriptions", "labevents", "chartevents"], + ... ) + >>> task = SepsisPredictionMIMIC4() + >>> samples = dataset.set_task(task) + """ + + task_name: str = "SepsisPredictionMIMIC4" + input_schema: ClassVar[dict[str, str]] = {"observations": "timeseries"} + output_schema: ClassVar[dict[str, str]] = {"sepsis": "binary"} + + RESP_RATE_ITEMID: ClassVar[str] = "220210" + SBP_ITEMIDS: ClassVar[list[str]] = ["220179", "220050"] + GCS_ITEMIDS: ClassVar[list[str]] = ["220739", "223900", "223901"] + VITAL_ITEMIDS: ClassVar[list[str]] = [RESP_RATE_ITEMID] + SBP_ITEMIDS + GCS_ITEMIDS + + LAB_ITEMIDS: ClassVar[list[str]] = [ + "50813", # Lactate + "50912", # Creatinine + "51265", # Platelet Count + "50885", # Bilirubin, Total + "51301", # White Blood Cells + "50882", # Bicarbonate + ] + + OBSERVATION_ITEMIDS: ClassVar[list[str]] = LAB_ITEMIDS + VITAL_ITEMIDS + + ANTIBIOTIC_DRUG_NAMES: ClassVar[list[str]] = [ + "vancomycin", + "cefepime", + "piperacillin", + "zosyn", + "meropenem", + "ceftriaxone", + "levofloxacin", + "ciprofloxacin", + "metronidazole", + "azithromycin", + "ampicillin", + "gentamicin", + "clindamycin", + "daptomycin", + "linezolid", + "imipenem", + ] + + ANTIBIOTIC_WINDOW_HOURS: ClassVar[float] = 24.0 + QSOFA_THRESHOLD: ClassVar[int] = 2 + GCS_NORMAL: ClassVar[int] = 15 + + def _pivot( + self, df: pl.DataFrame, table: str, item_ids: list[str] + ) -> pl.DataFrame: + """Pivot a filtered events frame to a wide timeseries matrix. + + Matches the pivot pattern used by ``InHospitalMortalityMIMIC4``: + one row per timestamp, one column per item ID, with missing item + IDs added as all-null columns so every returned frame has the same + shape regardless of which items were actually observed. + + Args: + df: A ``return_df=True`` events frame for a single table, + already time/``hadm_id``-filtered. + table: The source event type (e.g. ``"labevents"``), used to + resolve the ``{table}/itemid`` and ``{table}/valuenum`` + column names. + item_ids: The item IDs to keep as columns, in output order. + + Returns: + A DataFrame with a ``timestamp`` column plus one float column + per entry in ``item_ids``. + + Example: + Called from ``__call__`` to turn a raw ``chartevents`` or + ``labevents`` slice into the wide matrix ``_qsofa_onset`` and + the final ``observations`` tensor both expect, e.g. + ``self._pivot(vitals_df, "chartevents", self.VITAL_ITEMIDS)``. + """ + empty = pl.DataFrame({"timestamp": []}).cast({"timestamp": pl.Datetime}) + if df.height == 0: + return empty.with_columns([pl.lit(None).alias(i) for i in item_ids]) + + df = df.filter(pl.col(f"{table}/itemid").is_in(item_ids)) + if df.height == 0: + return empty.with_columns([pl.lit(None).alias(i) for i in item_ids]) + + df = df.select( + pl.col("timestamp"), + pl.col(f"{table}/itemid"), + pl.col(f"{table}/valuenum").cast(pl.Float64), + ) + df = df.pivot( + index="timestamp", + on=f"{table}/itemid", + values=f"{table}/valuenum", + aggregate_function="first", + ) + missing = [i for i in item_ids if i not in df.columns] + for col in missing: + df = df.with_columns(pl.lit(None).alias(col)) + return df.select("timestamp", *item_ids) + + def _qsofa_onset(self, vitals: pl.DataFrame) -> datetime | None: + """Find the earliest timestamp where qSOFA crosses the threshold. + + At each vitals timestamp, scores 1 point each for respiratory + rate >= 22, systolic BP <= 100 (the minimum across whichever BP + item IDs are present at that timestamp), and altered mental status + (summed GCS eye/verbal/motor < ``GCS_NORMAL``, scored only when + all three components are present). + + Args: + vitals: The pivoted vitals frame returned by ``_pivot``, with + one row per timestamp and one column per vital item ID. + + Returns: + The first timestamp where the qSOFA point total reaches + ``QSOFA_THRESHOLD``, or ``None`` if it never does. + + Example: + Called from ``__call__`` on the pivoted full-admission vitals + to find the onset timestamp used both as the sepsis-label + trigger and as the feature-censoring cutoff, e.g. + ``self._qsofa_onset(full_vitals)``. + """ + vitals = vitals.sort("timestamp") + for row in vitals.iter_rows(named=True): + rr = row[self.RESP_RATE_ITEMID] + sbp_candidates = [ + row[i] for i in self.SBP_ITEMIDS if row[i] is not None + ] + sbp = min(sbp_candidates) if sbp_candidates else None + gcs_components = [ + row[i] for i in self.GCS_ITEMIDS if row[i] is not None + ] + gcs_total = sum(gcs_components) if len(gcs_components) == 3 else None + + points = 0 + if rr is not None and rr >= 22: + points += 1 + if sbp is not None and sbp <= 100: + points += 1 + if gcs_total is not None and gcs_total < self.GCS_NORMAL: + points += 1 + + if points >= self.QSOFA_THRESHOLD: + return row["timestamp"] + return None + + def _antibiotic_times( + self, patient: Any, hadm_id: Any, start: datetime, end: datetime + ) -> list[datetime]: + """Return timestamps of antibiotic orders during an admission. + + Filters the admission's prescriptions to rows whose free-text + ``drug`` name contains any entry in ``ANTIBIOTIC_DRUG_NAMES`` + (case-insensitive substring match) -- see the class docstring's + Limitations section for why this is a heuristic, not a guarantee. + + Args: + patient: The patient to query. + hadm_id: The admission ID to restrict prescriptions to. + start: Window start (inclusive). + end: Window end (inclusive). + + Returns: + Timestamps of matching antibiotic orders, in no particular + order. + + Example: + Called from ``__call__`` once a qSOFA onset time is found, to + check whether an antibiotic order exists nearby, e.g. + ``self._antibiotic_times(patient, admission.hadm_id, + admission.timestamp, dischtime)``. + """ + prescriptions = patient.get_events( + event_type="prescriptions", + start=start, + end=end, + filters=[("hadm_id", "==", hadm_id)], + return_df=True, + ) + if prescriptions.height == 0: + return [] + pattern = "|".join(self.ANTIBIOTIC_DRUG_NAMES) + prescriptions = prescriptions.filter( + pl.col("prescriptions/drug").str.to_lowercase().str.contains(pattern) + ) + return prescriptions["timestamp"].to_list() + + def __call__(self, patient: Any) -> list[dict[str, Any]]: + """Build one sepsis-prediction sample per completed admission. + + For each admission: determines the earliest qSOFA-positive + timestamp (if any) across the full admission window, then checks + for an antibiotic order within ``ANTIBIOTIC_WINDOW_HOURS`` of it. + If both signals are found, the admission is labeled ``sepsis=1`` + with features censored strictly before that onset time. Otherwise + it is labeled ``sepsis=0`` using the full admission window as + features. Admissions with a missing/unparseable discharge time, or + with no observations remaining after censoring, are skipped. + + Args: + patient: The patient whose admissions to process. + + Returns: + A list of sample dicts, each with ``patient_id``, + ``admission_id``, ``observations`` (a ``(timestamps, values)`` + tuple matching the ``"timeseries"`` processor's input + contract), and ``sepsis``. + + Example: + Never called directly by users -- invoked once per patient by + ``BaseDataset.set_task()``, e.g. + ``dataset.set_task(SepsisPredictionMIMIC4())``. See + ``tests/core/test_sepsis_prediction_mimic4.py`` for direct + usage against synthetic ``Patient`` objects, and the class + docstring above for the ``set_task`` example. + """ + samples: list[dict[str, Any]] = [] + admissions = patient.get_events(event_type="admissions") + + for admission in admissions: + try: + # MIMIC-IV timestamps are naive by design (de-identified, + # date-shifted, no timezone). + dischtime = datetime.strptime( # noqa: DTZ007 + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (TypeError, ValueError): + continue + if dischtime <= admission.timestamp: + continue + + full_vitals_df = patient.get_events( + event_type="chartevents", + start=admission.timestamp, + end=dischtime, + filters=[("hadm_id", "==", admission.hadm_id)], + return_df=True, + ) + full_vitals = self._pivot(full_vitals_df, "chartevents", self.VITAL_ITEMIDS) + onset_time = ( + self._qsofa_onset(full_vitals) if full_vitals.height > 0 else None + ) + + sepsis_label = 0 + cutoff = dischtime + if onset_time is not None: + antibiotic_times = self._antibiotic_times( + patient, admission.hadm_id, admission.timestamp, dischtime + ) + window = timedelta(hours=self.ANTIBIOTIC_WINDOW_HOURS) + if any(abs(onset_time - t) <= window for t in antibiotic_times): + sepsis_label = 1 + cutoff = onset_time + + if cutoff <= admission.timestamp: + continue + + # get_events(end=...) is inclusive, so the observation exactly at + # `cutoff` (the qSOFA-triggering row, when sepsis_label == 1) + # must be dropped explicitly -- otherwise the label-defining + # observation itself would leak into the model's input. + labs_df = patient.get_events( + event_type="labevents", + start=admission.timestamp, + end=cutoff, + filters=[("hadm_id", "==", admission.hadm_id)], + return_df=True, + ) + if labs_df.height > 0: + labs_df = labs_df.filter(pl.col("timestamp") < cutoff) + vitals_df = patient.get_events( + event_type="chartevents", + start=admission.timestamp, + end=cutoff, + filters=[("hadm_id", "==", admission.hadm_id)], + return_df=True, + ) + if vitals_df.height > 0: + vitals_df = vitals_df.filter(pl.col("timestamp") < cutoff) + labs = self._pivot(labs_df, "labevents", self.LAB_ITEMIDS) + vitals = self._pivot(vitals_df, "chartevents", self.VITAL_ITEMIDS) + + observations = labs.join(vitals, on="timestamp", how="full", coalesce=True) + if observations.height == 0: + continue + observations = observations.sort("timestamp") + observations = observations.select("timestamp", *self.OBSERVATION_ITEMIDS) + + timestamps = observations["timestamp"].to_list() + values = observations.drop("timestamp").to_numpy() + + samples.append( + { + "patient_id": patient.patient_id, + "admission_id": admission.hadm_id, + "observations": (timestamps, values), + "sepsis": sepsis_label, + } + ) + + return samples diff --git a/tests/core/test_sepsis_prediction_mimic4.py b/tests/core/test_sepsis_prediction_mimic4.py new file mode 100644 index 000000000..16c23ac47 --- /dev/null +++ b/tests/core/test_sepsis_prediction_mimic4.py @@ -0,0 +1,430 @@ +# Author: Anish Gupta +# NetID: anishg8 +# Paper Title: N/A (original task contribution, not a paper reproduction) +# Paper Link: N/A +# Description: Synthetic-data test suite for SepsisPredictionMIMIC4. All +# fixtures are hand-built, in-memory Patient objects (2-8 fake patients, +# a handful of events each) -- no real MIMIC-IV files, no network, no +# downloads. The full suite runs in well under a second. +"""Tests for SepsisPredictionMIMIC4. + +The task is applied directly to hand-built ``Patient`` objects (bypassing +``BaseDataset``/YAML loading entirely), matching the layout PyHealth's +``global_event_df`` uses internally: one wide, sparse polars DataFrame with +an ``event_type``/``timestamp`` pair per row and ``{event_type}/{attr}`` +prefixed attribute columns. This keeps the fixtures small and avoids needing +real MIMIC-IV files (including the ``d_items.csv.gz``/``d_labitems.csv.gz`` +lookup joins) just to exercise the task's own logic. + +Manual test cases +------------------ +The automated tests below double as the manually-specified test cases +required by the contribution guide. Each one is stated here as +input -> expected output for quick review, independent of reading the code: + +1. ``test_qsofa_and_antibiotic_together_trigger_sepsis`` + Input: RR=25 + SBP=90 at hour 6 (qSOFA=2), antibiotic order at hour 5 + (within the 24h window), plus normal vitals at hour 1 and a lab at + hour 3 and hour 8. + Expected: 1 sample, ``sepsis=1``, only pre-hour-6 observations present + (the hour-8 reading must not appear -- leakage guard). + +2. ``test_normal_vitals_never_trigger_sepsis`` + Input: RR/SBP always normal across two vitals timestamps + one lab. + Expected: 1 sample, ``sepsis=0``, all 3 observation timestamps retained + (no censoring for a negative label). + +3. ``test_altered_mentation_contributes_to_qsofa`` + Input: RR=24 (1 point) + GCS eye/verbal/motor summing to 8 < 15 (1 + point), SBP unset, antibiotic nearby. + Expected: 1 sample, ``sepsis=1`` -- proves the GCS-based criterion + alone can supply the second qSOFA point. + +4. ``test_sbp_uses_minimum_across_available_itemids`` + Input: RR=24 (1 point), NIBP SBP=120 (normal) and arterial SBP=85 (low) + at the same timestamp, antibiotic nearby. + Expected: 1 sample, ``sepsis=1`` -- proves the lower of the two BP + readings is the one used, not whichever happens to be first. + +5. ``test_multiple_admissions_labeled_independently`` + Input: one patient with two admissions 30 days apart -- the first + qualifies for sepsis, the second has normal vitals. + Expected: 2 samples, one per admission, each labeled independently + (``sepsis=1`` and ``sepsis=0`` respectively). + +6. ``test_onset_with_no_prior_observations_yields_no_sample`` + Input: the only vitals for the admission are the exact qSOFA-triggering + reading itself. + Expected: 0 samples -- after censoring strictly before onset, there is + nothing left to predict from, so the admission is dropped rather than + returned with an empty feature matrix. + +7. ``test_organ_dysfunction_without_infection_is_not_sepsis`` + Input: qSOFA=2 at hour 4, but the only antibiotic order is 5 days + later (outside the admission and the window). + Expected: 1 sample, ``sepsis=0`` -- proves organ dysfunction alone, + without a nearby infection signal, is not labeled sepsis. + +8. ``test_missing_dischtime_admission_is_skipped`` + Input: an admission whose ``dischtime`` is ``None``. + Expected: 0 samples -- malformed data is skipped, not raised. + +9. ``test_samples_flow_through_the_real_processors`` + Input: one positive and one negative sample, run through + ``create_sample_dataset`` with the task's real ``input_schema``/ + ``output_schema``. + Expected: both samples processed successfully; the positive sample's + ``observations`` tensor has width ``len(OBSERVATION_ITEMIDS)`` and its + ``sepsis`` value is 1.0. + +10. ``test_tiny_model_trains_one_step`` + Input: the same two samples as above, fed through an ``RNN`` + (``hidden_dim=4``) and one ``Trainer.train()`` epoch + (``batch_size=2``). + Expected: training completes without error and produces a finite, + non-NaN loss -- confirms the task's output is actually consumable by + a real model, not just schema-shaped. +""" + +import unittest +from datetime import datetime + +import numpy as np +import polars as pl +import torch + +from pyhealth.data import Patient +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import RNN +from pyhealth.tasks.sepsis_prediction_mimic4 import SepsisPredictionMIMIC4 +from pyhealth.trainer import Trainer + +T0 = datetime(2024, 1, 1, 0, 0, 0) # naive, matching MIMIC-IV's own timestamps # noqa: DTZ001 + + +def _row(event_type: str, timestamp: datetime, **attrs) -> dict: + row = {"event_type": event_type, "timestamp": timestamp} + for key, value in attrs.items(): + row[f"{event_type}/{key}"] = value + return row + + +# A real PyHealth dataset guarantees every configured table's attribute +# columns exist in the global event frame, regardless of whether any given +# admission has rows for that table -- get_events() filters against +# "{event_type}/{attr}" columns unconditionally. These fixtures must +# guarantee the same columns exist even when a scenario has zero rows for +# a given event type (e.g. no prescriptions at all). +_REQUIRED_COLUMNS = { + "admissions/hadm_id": pl.Int64, + "admissions/dischtime": pl.Utf8, + "prescriptions/hadm_id": pl.Int64, + "prescriptions/drug": pl.Utf8, + "labevents/hadm_id": pl.Int64, + "labevents/itemid": pl.Utf8, + "labevents/valuenum": pl.Float64, + "chartevents/hadm_id": pl.Int64, + "chartevents/itemid": pl.Utf8, + "chartevents/valuenum": pl.Float64, +} + + +def _build_patient(patient_id: str, rows: list) -> Patient: + df = pl.DataFrame(rows) + for col, dtype in _REQUIRED_COLUMNS.items(): + if col not in df.columns: + df = df.with_columns(pl.lit(None, dtype=dtype).alias(col)) + df = df.cast(_REQUIRED_COLUMNS) + return Patient(patient_id=patient_id, data_source=df) + + +def _admission(hadm_id: int, admit: datetime, dischtime: datetime) -> dict: + return _row( + "admissions", + admit, + hadm_id=hadm_id, + dischtime=dischtime.strftime("%Y-%m-%d %H:%M:%S"), + ) + + +def _vital(hadm_id: int, ts: datetime, itemid: str, valuenum: float) -> dict: + return _row( + "chartevents", ts, hadm_id=hadm_id, itemid=itemid, valuenum=valuenum + ) + + +def _lab(hadm_id: int, ts: datetime, itemid: str, valuenum: float) -> dict: + return _row( + "labevents", ts, hadm_id=hadm_id, itemid=itemid, valuenum=valuenum + ) + + +def _antibiotic(hadm_id: int, ts: datetime, drug: str = "Vancomycin") -> dict: + return _row("prescriptions", ts, hadm_id=hadm_id, drug=drug) + + +class TestSepsisPredictionMIMIC4(unittest.TestCase): + def setUp(self): + self.task = SepsisPredictionMIMIC4() + + def test_qsofa_and_antibiotic_together_trigger_sepsis(self): + from datetime import timedelta + + rows = [ + _admission(1001, T0, T0 + timedelta(hours=10)), + # Normal vitals early in the stay. + _vital(1001, T0 + timedelta(hours=1), "220210", 16.0), # RR normal + _vital(1001, T0 + timedelta(hours=1), "220179", 120.0), # SBP normal + # A relevant lab before onset. + _lab(1001, T0 + timedelta(hours=3), "50813", 4.5), # Lactate + # qSOFA >= 2 at hour 6: RR high + SBP low. + _vital(1001, T0 + timedelta(hours=6), "220210", 25.0), + _vital(1001, T0 + timedelta(hours=6), "220179", 90.0), + # Antibiotic within the 24h window of onset. + _antibiotic(1001, T0 + timedelta(hours=5)), + # A later vital that must NOT leak into the returned features. + _vital(1001, T0 + timedelta(hours=8), "220210", 30.0), + ] + patient = _build_patient("p1", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertEqual(sample["sepsis"], 1) + + timestamps, values = sample["observations"] + self.assertTrue(all(t < T0 + timedelta(hours=6) for t in timestamps)) + # The hour-8 RR=30.0 value must not appear anywhere in the matrix. + self.assertNotIn(30.0, values.flatten().tolist()) + + def test_normal_vitals_never_trigger_sepsis(self): + from datetime import timedelta + + rows = [ + _admission(1002, T0, T0 + timedelta(hours=8)), + _vital(1002, T0 + timedelta(hours=1), "220210", 14.0), + _vital(1002, T0 + timedelta(hours=1), "220179", 118.0), + _vital(1002, T0 + timedelta(hours=5), "220210", 15.0), + _vital(1002, T0 + timedelta(hours=5), "220179", 122.0), + _lab(1002, T0 + timedelta(hours=2), "50912", 0.9), # Creatinine + ] + patient = _build_patient("p2", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["sepsis"], 0) + timestamps, _ = samples[0]["observations"] + # 3 distinct timestamps: hour-1 vitals, hour-2 lab, hour-5 vitals + # (the RR+SBP pair at each vitals timestamp collapse into one row). + self.assertEqual(len(timestamps), 3) + + def test_altered_mentation_contributes_to_qsofa(self): + """The GCS/altered-mentation criterion must itself be able to + contribute a qSOFA point -- every other test only exercises RR and + SBP, leaving the GCS-scoring branch entirely untested.""" + from datetime import timedelta + + rows = [ + _admission(1004, T0, T0 + timedelta(hours=10)), + # A baseline reading, so censoring at onset leaves something. + _vital(1004, T0 + timedelta(hours=1), "220210", 16.0), + # RR alone (1 point) + GCS sum=8 < 15 (1 point) = 2 points. + # SBP is left unset (no reading), so it cannot contribute. + _vital(1004, T0 + timedelta(hours=4), "220210", 24.0), + _vital(1004, T0 + timedelta(hours=4), "220739", 2.0), # eye + _vital(1004, T0 + timedelta(hours=4), "223900", 2.0), # verbal + _vital(1004, T0 + timedelta(hours=4), "223901", 4.0), # motor + _antibiotic(1004, T0 + timedelta(hours=4)), + ] + patient = _build_patient("p4", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["sepsis"], 1) + + def test_sbp_uses_minimum_across_available_itemids(self): + """When both an NIBP and an arterial-line SBP reading exist at the + same timestamp, the lower of the two must be used -- proving the + min() selection, not just picking whichever itemid happens first.""" + from datetime import timedelta + + rows = [ + _admission(1005, T0, T0 + timedelta(hours=10)), + # A baseline reading, so censoring at onset leaves something. + _vital(1005, T0 + timedelta(hours=1), "220179", 122.0), + # NIBP reads normal (120); arterial line reads low (85). + # Only the minimum (85 <= 100) should count as a qSOFA point. + _vital(1005, T0 + timedelta(hours=3), "220210", 24.0), # RR point + _vital(1005, T0 + timedelta(hours=3), "220179", 120.0), # NIBP + _vital(1005, T0 + timedelta(hours=3), "220050", 85.0), # arterial + _antibiotic(1005, T0 + timedelta(hours=3)), + ] + patient = _build_patient("p5", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["sepsis"], 1) + + def test_multiple_admissions_labeled_independently(self): + """A patient with two admissions must get one sample per admission, + each labeled from its own vitals/antibiotics -- not pooled.""" + from datetime import timedelta + + first_admit = T0 + second_admit = T0 + timedelta(days=30) + rows = [ + _admission(2001, first_admit, first_admit + timedelta(hours=10)), + _vital(2001, first_admit + timedelta(hours=1), "220210", 16.0), + _vital(2001, first_admit + timedelta(hours=4), "220210", 25.0), + _vital(2001, first_admit + timedelta(hours=4), "220179", 90.0), + _antibiotic(2001, first_admit + timedelta(hours=4)), + _admission(2002, second_admit, second_admit + timedelta(hours=8)), + _vital(2002, second_admit + timedelta(hours=1), "220210", 14.0), + _vital(2002, second_admit + timedelta(hours=1), "220179", 118.0), + ] + patient = _build_patient("p6", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 2) + by_admission = {s["admission_id"]: s for s in samples} + self.assertEqual(by_admission[2001]["sepsis"], 1) + self.assertEqual(by_admission[2002]["sepsis"], 0) + + def test_onset_with_no_prior_observations_yields_no_sample(self): + """If the only vitals for an admission are the exact qSOFA- + triggering reading itself, strict pre-onset censoring leaves no + observations at all -- the sample must be dropped, not returned + with an empty feature matrix.""" + from datetime import timedelta + + rows = [ + _admission(1007, T0, T0 + timedelta(hours=10)), + _vital(1007, T0 + timedelta(hours=4), "220210", 25.0), + _vital(1007, T0 + timedelta(hours=4), "220179", 90.0), + _antibiotic(1007, T0 + timedelta(hours=4)), + ] + patient = _build_patient("p8", rows) + samples = self.task(patient) + + self.assertEqual(samples, []) + + def test_missing_dischtime_admission_is_skipped(self): + """An admission with a null/unparseable discharge time must be + skipped rather than raising.""" + rows = [ + _row( + "admissions", + T0, + hadm_id=1006, + dischtime=None, + ), + ] + patient = _build_patient("p7", rows) + samples = self.task(patient) + + self.assertEqual(samples, []) + + def test_organ_dysfunction_without_infection_is_not_sepsis(self): + """qSOFA >= 2 alone, with no antibiotic order anywhere nearby, must + not be labeled sepsis -- this is what proves the task enforces both + signals rather than just a vitals threshold.""" + from datetime import timedelta + + rows = [ + _admission(1003, T0, T0 + timedelta(hours=10)), + _vital(1003, T0 + timedelta(hours=4), "220210", 25.0), + _vital(1003, T0 + timedelta(hours=4), "220179", 90.0), + # An antibiotic order, but far outside the 24h onset window. + _antibiotic(1003, T0 + timedelta(days=5)), + ] + patient = _build_patient("p3", rows) + samples = self.task(patient) + + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["sepsis"], 0) + + def test_samples_flow_through_the_real_processors(self): + """End-to-end: raw task output must be accepted by the actual + 'timeseries'/'binary' processors used by set_task(), not just be a + plausible-looking dict.""" + from datetime import timedelta + + positive_rows = [ + _admission(1001, T0, T0 + timedelta(hours=10)), + _vital(1001, T0 + timedelta(hours=1), "220210", 16.0), + _vital(1001, T0 + timedelta(hours=1), "220179", 120.0), + _vital(1001, T0 + timedelta(hours=6), "220210", 25.0), + _vital(1001, T0 + timedelta(hours=6), "220179", 90.0), + _antibiotic(1001, T0 + timedelta(hours=5)), + ] + negative_rows = [ + _admission(1002, T0, T0 + timedelta(hours=8)), + _vital(1002, T0 + timedelta(hours=1), "220210", 14.0), + _vital(1002, T0 + timedelta(hours=1), "220179", 118.0), + ] + samples = self.task(_build_patient("p1", positive_rows)) + samples += self.task(_build_patient("p2", negative_rows)) + # BinaryLabelProcessor.fit requires seeing both classes; fit on both + # samples, then assert on the positive one. + self.assertEqual(len(samples), 2) + + sample_dataset = create_sample_dataset( + samples=samples, + input_schema=self.task.input_schema, + output_schema=self.task.output_schema, + dataset_name="sepsis_test", + ) + self.assertEqual(len(sample_dataset), 2) + processed = next(s for s in sample_dataset if s["admission_id"] == 1001) + # (timesteps, num_observation_itemids) + self.assertEqual( + processed["observations"].shape[1], + len(SepsisPredictionMIMIC4.OBSERVATION_ITEMIDS), + ) + self.assertEqual(float(processed["sepsis"]), 1.0) + + def test_tiny_model_trains_one_step(self): + """Smoke test: task output must be trainable by a real model, not + just schema-shaped. Uses a tiny RNN (hidden_dim=4) and a single + epoch/batch so this stays a millisecond-scale unit test.""" + from datetime import timedelta + + torch.manual_seed(42) + np.random.seed(42) + + positive_rows = [ + _admission(1001, T0, T0 + timedelta(hours=10)), + _vital(1001, T0 + timedelta(hours=1), "220210", 16.0), + _vital(1001, T0 + timedelta(hours=6), "220210", 25.0), + _vital(1001, T0 + timedelta(hours=6), "220179", 90.0), + _antibiotic(1001, T0 + timedelta(hours=5)), + ] + negative_rows = [ + _admission(1002, T0, T0 + timedelta(hours=8)), + _vital(1002, T0 + timedelta(hours=1), "220210", 14.0), + _vital(1002, T0 + timedelta(hours=1), "220179", 118.0), + ] + samples = self.task(_build_patient("p1", positive_rows)) + samples += self.task(_build_patient("p2", negative_rows)) + self.assertEqual(len(samples), 2) + + sample_dataset = create_sample_dataset( + samples=samples, + input_schema=self.task.input_schema, + output_schema=self.task.output_schema, + dataset_name="sepsis_tiny_model_test", + ) + dataloader = get_dataloader(sample_dataset, batch_size=2, shuffle=False) + + model = RNN(dataset=sample_dataset, embedding_dim=4, hidden_dim=4) + trainer = Trainer(model=model, enable_logging=False) + trainer.train(train_dataloader=dataloader, epochs=1) + + batch = next(iter(dataloader)) + output = model(**batch) + loss = float(output["loss"]) + self.assertTrue(np.isfinite(loss)) + + +if __name__ == "__main__": + unittest.main()