diff --git a/.gitignore b/.gitignore index 086c8da3f..f0c1ff431 100644 --- a/.gitignore +++ b/.gitignore @@ -140,4 +140,18 @@ data/physionet.org/ .codex # Model weight files (large binaries, distributed separately) -weightfiles/ \ No newline at end of file +weightfiles/ + +# Local / personal (never commit) +CLAUDE.local.md +.claude/settings.local.json + +# Local test fixtures (download separately; not part of the repo) +test-resources/core/chestxray14/ +test-resources/meds_demo/ + +# Local Python environments (not .venv — that pattern is already ignored) +.venv312/ + +# Tool caches +.ruff_cache/ \ No newline at end of file diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 592aed487..c9a88b7ff 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -225,6 +225,7 @@ Available Datasets datasets/pyhealth.datasets.MIMIC3Dataset datasets/pyhealth.datasets.MIMIC4Dataset datasets/pyhealth.datasets.FHIRDataset + datasets/pyhealth.datasets.MEDSDataset datasets/pyhealth.datasets.MIMIC4FHIR datasets/pyhealth.datasets.MedicalTranscriptionsDataset datasets/pyhealth.datasets.CardiologyDataset diff --git a/docs/api/datasets/pyhealth.datasets.MEDSDataset.rst b/docs/api/datasets/pyhealth.datasets.MEDSDataset.rst new file mode 100644 index 000000000..92bc93e0d --- /dev/null +++ b/docs/api/datasets/pyhealth.datasets.MEDSDataset.rst @@ -0,0 +1,9 @@ +pyhealth.datasets.MEDSDataset +=================================== + +Dataset class for data in the `Medical Event Data Standard (MEDS) `_, a minimal event-based schema for machine learning over EHR data (MEDS Working Group / Arnrich et al., ICLR 2024 Workshop on Learning from Time Series For Health; `openreview:IsHy2ebjIG `_). Sharded Parquet event files are read with their native types, and standard MEDS splits (train / tuning / held_out) can be selected directly via the ``subset`` argument. The canonical subject-to-split mapping is defined in ``metadata/subject_splits.parquet``; see the `MEDS schema documentation `_. + +.. autoclass:: pyhealth.datasets.MEDSDataset + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index c7910e626..bdaa9599a 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -207,6 +207,7 @@ Available Tasks Base Task In-Hospital Mortality (MIMIC-IV) + In-Hospital Mortality (MEDS) MIMIC-III ICD-9 Coding Cardiology Detection COVID-19 CXR Classification diff --git a/docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst b/docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst new file mode 100644 index 000000000..9ca8c55d2 --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.InHospitalMortalityMEDS.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.InHospitalMortalityMEDS +====================================== + +.. autoclass:: pyhealth.tasks.in_hospital_mortality_meds.InHospitalMortalityMEDS + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/meds_demo.py b/examples/meds_demo.py new file mode 100644 index 000000000..faebaf78d --- /dev/null +++ b/examples/meds_demo.py @@ -0,0 +1,138 @@ +"""End-to-end example: loading a MEDS dataset with PyHealth. + +This example uses the public *MIMIC-IV demo data in the Medical Event Data +Standard (MEDS)* (PhysioNet, v0.0.1, ODbL v1.0, ~100 subjects): +https://doi.org/10.13026/t2y8-ea41 + +Download it once (open access, ~a few MB): + + wget -r -N -c -np https://physionet.org/files/mimic-iv-demo-meds/0.0.1/ + +Then run: + + python examples/meds_demo.py \\ + --root physionet.org/files/mimic-iv-demo-meds/0.0.1 + +Any dataset following the MEDS layout (``data/**.parquet`` + +``metadata/subject_splits.parquet``) works the same way. See the MEDS +specification: https://github.com/Medical-Event-Data-Standard/meds +""" + +import argparse +import tempfile +from pathlib import Path +from typing import Any, List + +import polars as pl +import pyhealth.datasets.configs as meds_configs + +from pyhealth.datasets import MEDSDataset, get_dataloader, split_by_patient +from pyhealth.models import RNN +from pyhealth.tasks import InHospitalMortalityMEDS +from pyhealth.trainer import Trainer + +# Full-stay MEDS code sequences are often thousands of events long. Keeping +# every code makes a vanilla RNN prohibitively slow on CPU, so the training +# block below keeps only the *most recent* codes per stay. This is a demo +# ergonomics choice, not a benchmark configuration: production runs should +# use the unmodified task (``InHospitalMortalityMEDS()``) and an appropriate +# model/window for the sequence lengths involved. +_DEMO_MAX_SEQ_LEN = 256 + + +class _DemoMortalityTask(InHospitalMortalityMEDS): + """Demo-only wrapper that tail-truncates ``codes`` before ``set_task``.""" + + def __init__(self, max_seq_len: int = _DEMO_MAX_SEQ_LEN, **kwargs) -> None: + super().__init__(**kwargs) + self.max_seq_len = max_seq_len + + def __call__(self, patient: Any) -> List[dict]: + samples = super().__call__(patient) + for sample in samples: + codes = sample["codes"] + if len(codes) > self.max_seq_len: + sample["codes"] = codes[-self.max_seq_len :] + return samples + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + required=True, + help="Root of the MEDS dataset (directory containing data/ and metadata/)", + ) + parser.add_argument( + "--subset", + default="train", + help="Split to load as a subset (default: train)", + ) + args = parser.parse_args() + + # 1) Load the full dataset: every Parquet shard under data/ is read, + # including nested split directories (data//.parquet). + dataset = MEDSDataset(root=args.root) + dataset.stats() + + # 2) Peek at the canonical event frame (typed straight from Parquet: + # string patient ids, datetime64[ms] timestamps, float values). + events = dataset.global_event_df + print(events.head(5).collect()) + + # 3) Load a split-restricted subset. Subjects are selected through the + # metadata/subject_splits.parquet assignment; each subset uses its + # own processing cache. + subset = MEDSDataset(root=args.root, subset=args.subset) + n_subset = len(subset.unique_patient_ids) + n_total = len(dataset.unique_patient_ids) + print(f"Subjects in subset '{args.subset}': {n_subset} / {n_total}") + + # 4) Static (null-time) MEDS events, e.g. demographics, are preserved. + n_static = ( + events.filter(pl.col("timestamp").is_null()).select(pl.len()).collect().item() + ) + print(f"Static (null-time) events: {n_static}") + + # 5) In-hospital mortality task + minimal RNN training loop (1 epoch). + # Sequences are tail-truncated via _DemoMortalityTask (see module note). + cfg = Path(meds_configs.__file__).parent / "meds_with_hadm.yaml" + task_cache = tempfile.mkdtemp(prefix="meds_demo_task_") + cohort = MEDSDataset( + root=args.root, + config_path=str(cfg), + subset=args.subset, + cache_dir=task_cache, + ) + samples = cohort.set_task(_DemoMortalityTask()) + print( + f"Mortality task samples ({args.subset}, codes tail-truncated to " + f"{_DEMO_MAX_SEQ_LEN}): {len(samples)}" + ) + + train_dataset, val_dataset, test_dataset = split_by_patient( + samples, [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=samples, embedding_dim=64, hidden_dim=64) + trainer = Trainer(model=model) + trainer.train( + train_dataloader=train_dataloader, + val_dataloader=val_dataloader, + epochs=1, + monitor="roc_auc", + ) + metrics = trainer.evaluate(test_dataloader) + print( + "Test metrics (smoke test only — 1 epoch, tail-truncated sequences; " + "not interpretable as model quality):" + ) + print(metrics) + + +if __name__ == "__main__": + # BaseDataset spawns Dask worker processes; keep the main-module guard. + main() diff --git a/examples/verify_meds_mortality.py b/examples/verify_meds_mortality.py new file mode 100644 index 000000000..8c199db84 --- /dev/null +++ b/examples/verify_meds_mortality.py @@ -0,0 +1,62 @@ +"""Standalone verification of the MEDS in-hospital mortality cohort. + +Applies ``InHospitalMortalityMEDS`` via ``set_task`` on a MEDS dataset and +prints basic cohort counts (expected ~12/238 positives on the public +MIMIC-IV demo). Prefer this path over a per-patient ``task(get_patient)`` +loop: ``set_task`` uses the library's lazy loading and parallel processing. + +Usage: + # Download the public demo once (open access, ODbL v1.0): + # wget -r -N -c -np https://physionet.org/files/mimic-iv-demo-meds/0.0.1/ + python examples/verify_meds_mortality.py \\ + --root physionet.org/files/mimic-iv-demo-meds/0.0.1 + +The task needs hadm_id; this script uses the bundled +``configs/meds_with_hadm.yaml`` automatically. +""" + +import argparse +from pathlib import Path + +import pyhealth.datasets.configs as meds_configs +from pyhealth.datasets import MEDSDataset +from pyhealth.tasks import InHospitalMortalityMEDS + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + required=True, + help="Root of the MEDS dataset (contains data/ and metadata/).", + ) + parser.add_argument( + "--observation-window", + default="full_stay", + choices=["full_stay", "first_hours"], + ) + parser.add_argument("--window-hours", type=float, default=48.0) + args = parser.parse_args() + + cfg = Path(meds_configs.__file__).parent / "meds_with_hadm.yaml" + dataset = MEDSDataset(root=args.root, config_path=str(cfg)) + task = InHospitalMortalityMEDS( + observation_window=args.observation_window, + window_hours=args.window_hours, + ) + + samples = dataset.set_task(task) + n = len(samples) + n_positive = sum(int(samples[i]["mortality"]) for i in range(n)) + n_patients = len({samples[i]["patient_id"] for i in range(n)}) + + print(f"root : {args.root}") + print(f"observation_window : {args.observation_window}") + print(f"n_samples (stays) : {n}") + print(f"n_patients : {n_patients}") + print(f"n_positive (died) : {n_positive}") + print(f"positive_rate : {((n_positive / n) if n else 0.0):.4f}") + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 99d90aa62..d09b05e3d 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -57,6 +57,7 @@ def __init__(self, *args, **kwargs): from .eicu import eICUDataset from .isruc import ISRUCDataset from .medical_transcriptions import MedicalTranscriptionsDataset +from .meds import MEDSDataset as MEDSDataset # noqa: E402 from .mimic3 import MIMIC3Dataset from .mimic4 import MIMIC4CXRDataset, MIMIC4Dataset, MIMIC4EHRDataset, MIMIC4NoteDataset from .fhir import FHIRDataset, MIMIC4FHIR diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 0e4280aab..3d449d579 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -315,6 +315,15 @@ class BaseDataset(ABC): config (dict): Configuration loaded from a YAML file. global_event_df (pl.LazyFrame): The global event data frame. dev (bool): Whether to enable dev mode (limit to 1000 patients). + + Examples: + >>> from pyhealth.datasets import BaseDataset + >>> dataset = BaseDataset( + ... root="/path/to/source", + ... tables=["patients", "diagnoses"], + ... config_path="/path/to/config.yaml", + ... ) + >>> dataset.stats() """ def __init__( @@ -425,6 +434,68 @@ def clean_tmpdir(self) -> None: if tmp_dir.exists(): shutil.rmtree(tmp_dir) + def _scan_table(self, source_path: str) -> dd.DataFrame: + """Routes a table source to the appropriate scanner based on its format. + + Parquet sources (``.parquet``/``.pq`` files, glob patterns targeting + such files, or directories of Parquet shards) are handled by + :meth:`_scan_parquet`. Any other source falls back to the existing + CSV/TSV(.gz) scanner, preserving prior behavior for all datasets. + + Args: + source_path (str): Path to the table source. + + Returns: + dd.DataFrame: The Dask DataFrame for the table source. + """ + stripped = source_path.rstrip("/") + if stripped.endswith((".parquet", ".pq")) or ( + not is_url(source_path) and Path(source_path).is_dir() + ): + return self._scan_parquet(source_path) + return self._scan_csv_tsv_gz(source_path) + + def _scan_parquet(self, source_path: str) -> dd.DataFrame: + """Scans a Parquet source and returns a Dask DataFrame. + + The source may be a single ``.parquet``/``.pq`` file, a glob pattern, + or a directory that is scanned recursively — which supports sharded + datasets such as MEDS, laid out as ``data//.parquet``. + + Unlike :meth:`_scan_csv_tsv_gz`, no all-string schema coercion is + applied: Parquet files embed their schema, so source dtypes (native + timestamps, numeric columns, nullable strings) are preserved and + handled downstream by :meth:`load_table`. + + Args: + source_path (str): Path to a Parquet file, directory, or glob. + + Returns: + dd.DataFrame: The Dask DataFrame backed by the Parquet source. + + Raises: + FileNotFoundError: If the source path does not exist, or if a + directory source contains no Parquet files. + """ + path = Path(source_path) + is_glob = any(ch in source_path for ch in "*?[") + if not is_glob: + if not path.exists(): + raise FileNotFoundError( + f"Parquet source does not exist: {source_path}" + ) + if path.is_dir() and not any( + itertools.chain(path.rglob("*.parquet"), path.rglob("*.pq")) + ): + raise FileNotFoundError( + f"Directory contains no Parquet files: {source_path}" + ) + return dd.read_parquet( + source_path, + split_row_groups=True, # type: ignore + blocksize="64MB", + ) + def _scan_csv_tsv_gz(self, source_path: str) -> dd.DataFrame: """Scans a CSV/TSV file (possibly gzipped) and returns a Dask DataFrame. @@ -596,7 +667,8 @@ def load_table(self, table_name: str) -> dd.DataFrame: Raises: ValueError: If the table is not found in the config. - FileNotFoundError: If the CSV file for the table or join is not found. + FileNotFoundError: If the source file (CSV/TSV or Parquet) for the + table or join is not found. """ assert self.config is not None, "Config must be provided to load tables" @@ -608,7 +680,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: csv_path = clean_path(csv_path) logger.info(f"Scanning table: {table_name} from {csv_path}") - df = self._scan_csv_tsv_gz(csv_path) + df = self._scan_table(csv_path) # Convert column names to lowercase before calling preprocess_func df = df.rename(columns=str.lower) @@ -627,7 +699,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: other_csv_path = f"{self.root}/{join_cfg.file_path}" other_csv_path = clean_path(other_csv_path) logger.info(f"Joining with table: {other_csv_path}") - join_df = self._scan_csv_tsv_gz(other_csv_path) + join_df = self._scan_table(other_csv_path) join_df = join_df.rename(columns=str.lower) join_key = join_cfg.on columns = join_cfg.columns @@ -651,14 +723,21 @@ def load_table(self, table_name: str) -> dd.DataFrame: timestamp_series: dd.Series = functools.reduce( operator.add, (df[col].astype("string") for col in timestamp_col) ) + timestamp_series = dd.to_datetime( + timestamp_series, + format=timestamp_format, + errors="raise", + ) + elif pd.api.types.is_datetime64_any_dtype(df[timestamp_col].dtype): + # Typed sources (e.g. Parquet) already carry native timestamps: + # skip the string round-trip and only normalize the unit below. + timestamp_series: dd.Series = df[timestamp_col] else: - timestamp_series: dd.Series = df[timestamp_col].astype("string") - - timestamp_series: dd.Series = dd.to_datetime( - timestamp_series, - format=timestamp_format, - errors="raise", - ) + timestamp_series = dd.to_datetime( + df[timestamp_col].astype("string"), + format=timestamp_format, + errors="raise", + ) df: dd.DataFrame = df.assign( timestamp=timestamp_series.astype("datetime64[ms]") ) diff --git a/pyhealth/datasets/configs/meds.yaml b/pyhealth/datasets/configs/meds.yaml new file mode 100644 index 000000000..a251854ee --- /dev/null +++ b/pyhealth/datasets/configs/meds.yaml @@ -0,0 +1,34 @@ +# MEDS (Medical Event Data Standard) tables. +# +# MEDS is already flat (one row per measurement), so a single event table +# covers the data shards. The canonical subject-to-split mapping is exposed +# as a second, ordinary event table -- the same pattern as the `splits` +# table in ehrshot.yaml. TableConfig has no format field: Parquet reading is +# handled by BaseDataset._scan_table / _scan_parquet, so no config-schema +# change is needed for the 21 existing datasets. +# +# Paths match mimic-iv-demo-meds: data//*.parquet + +# metadata/subject_splits.parquet. +version: "1.0" +tables: + meds: + # Directory of shards; dd.read_parquet reads the nested tree whole. + file_path: "data" + patient_id: "subject_id" + # Native datetime64[us] in MEDS Parquet, narrowed to [ms] by the + # BaseDataset typed-timestamp fast-path. No timestamp_format: that + # would imply text parsing. MEDSDataset rejects non-timestamp / + # tz-aware columns at construction (footer schema guard). + timestamp: "time" + attributes: + - "code" + - "numeric_value" + + # Canonical split map as events (attribute `subject_splits/split`). + # Opt-in: tables=["meds", "subject_splits"]. + subject_splits: + file_path: "metadata/subject_splits.parquet" + patient_id: "subject_id" + timestamp: null + attributes: + - "split" diff --git a/pyhealth/datasets/configs/meds_with_hadm.yaml b/pyhealth/datasets/configs/meds_with_hadm.yaml new file mode 100644 index 000000000..508e901bc --- /dev/null +++ b/pyhealth/datasets/configs/meds_with_hadm.yaml @@ -0,0 +1,22 @@ +version: "1.0" +# Stay-aware MEDS config: identical to the default configs/meds.yaml but +# additionally exposes `hadm_id`, which stay-based tasks (e.g. +# InHospitalMortalityMEDS) need to reconstruct admissions/discharges. +# The default config keeps `attributes` minimal; opt into this one when a +# task consumes hadm_id. +tables: + meds: + file_path: "data" + patient_id: "subject_id" + timestamp: "time" + attributes: + - "code" + - "numeric_value" + - "hadm_id" + + subject_splits: + file_path: "metadata/subject_splits.parquet" + patient_id: "subject_id" + timestamp: null + attributes: + - "split" diff --git a/pyhealth/datasets/meds.py b/pyhealth/datasets/meds.py new file mode 100644 index 000000000..9a090e239 --- /dev/null +++ b/pyhealth/datasets/meds.py @@ -0,0 +1,305 @@ +"""MEDS (Medical Event Data Standard) dataset for PyHealth. + +MEDS distributes event data as *typed*, sharded Parquet, already flattened to +one row per measurement -- ``(subject_id, time, code, numeric_value, ...)`` -- +plus a canonical subject-to-split mapping at +``metadata/subject_splits.parquet``. See the MEDS schema documentation for +the canonical subject-to-split mapping: +https://medical-event-data-standard.github.io/ +This maps almost one-to-one onto +PyHealth's canonical event schema +(``patient_id | event_type | timestamp | /``). + +Parquet scanning and the typed-timestamp fast-path live in +:class:`BaseDataset`. ``MEDSDataset`` adds three MEDS-specific pieces: + +1. **Schema contract at construction.** :meth:`_validate_event_schema` reads + Parquet footers only and raises ``TypeError`` when the configured + timestamp column is missing, not a timestamp type, or timezone-aware + (MEDS reference ``DataSchema`` is ``timestamp[us]``, tz-naive). +2. **Split-aware loading.** ``subset=`` keeps only the patients of one + canonical split, via ``split_source`` (``"metadata"`` or ``"directory"``). +3. **Cache disambiguation.** Subset instances nest a dedicated cache + directory so different splits never share a processing cache. + +MEDS spec: https://github.com/Medical-Event-Data-Standard/meds +""" + +import logging +from pathlib import Path +from typing import List, Literal, Optional + +import dask.dataframe as dd +import pandas as pd +import pyarrow as pa +import pyarrow.dataset as pa_ds + +from .base_dataset import BaseDataset, clean_path + +logger = logging.getLogger(__name__) + +#: Canonical MEDS split names, in canonical order (MEDS spec). +MEDS_SPLITS: tuple[str, ...] = ("train", "tuning", "held_out") + +#: MEDS-normative locations, relative to the dataset root. +DATA_RELPATH = "data" +SUBJECT_SPLITS_RELPATH = "metadata/subject_splits.parquet" + +SplitSource = Literal["metadata", "directory"] + + +class MEDSDataset(BaseDataset): + """Dataset for MEDS (Medical Event Data Standard) sources. + + MEDS data is distributed as sharded, typed Parquet under per-split + directories (``data/train/*.parquet``, ``data/tuning/*.parquet``, + ``data/held_out/*.parquet``) plus a canonical subject-to-split map at + ``metadata/subject_splits.parquet``. See the MEDS schema documentation: + https://medical-event-data-standard.github.io/ + + ``time`` must be a timezone-naive timestamp (MEDS reference schema); + violations raise ``TypeError`` at construction. + + Split handling: + The canonical split is available two ways, both optional: + + * as **events**: load the ``subject_splits`` table + (``tables=["meds", "subject_splits"]``) and each subject carries one + ``subject_splits`` event with attribute ``subject_splits/split`` -- + the exact pattern of EHRShot's ``splits`` table, usable from + ``Task.pre_filter`` or per-patient logic; + * as a **loader filter**: ``subset="train"`` (or ``"tuning"`` / + ``"held_out"``) keeps only that split's patients in every loaded + table, via the same patient-``isin`` mechanic as dev mode. + + ``split_source`` controls where ``subset`` gets its patient list: + ``"metadata"`` (default) reads the canonical mapping file -- + authoritative per the MEDS spec and independent of directory layout; + ``"directory"`` derives it from which ``data//`` directory + subjects appear in -- useful when an export omits the metadata file. + The two sources *should* agree; whether PyHealth must verify that + equivalence is an open question for the upstream maintainer, so + this class does not silently pick one when they could diverge: it + uses exactly the source you asked for, and caches them separately. + + Note: + ``event_type`` is the table name (``"meds"``) for every row; the + clinically meaningful event kind lives in the ``meds/code`` + attribute. This mirrors EHRShot, whose single ``ehrshot`` table also + carries an event vocabulary in a ``code`` attribute. Whether upstream + prefers mapping MEDS ``code`` onto ``event_type`` instead is a design + question for the maintainer. + + Args: + root: Root directory of the MEDS dataset (the directory that + contains ``data/`` and ``metadata/``). + tables: Tables to load, as named in ``configs/meds.yaml``. Defaults + to ``["meds"]``; add ``"subject_splits"`` to expose the canonical + split as events. + subset: ``"train"``, ``"tuning"``, ``"held_out"``, or ``"all"`` + (default). Anything but ``"all"`` filters every loaded table to + that split's patients. + split_source: Where ``subset`` gets its patient list from; see + above. Ignored when ``subset="all"``. + dataset_name: Dataset name. Defaults to ``"meds"``. + config_path: Path to the YAML config. Defaults to + ``configs/meds.yaml``. + **kwargs: Forwarded to :class:`BaseDataset` (``cache_dir``, + ``num_workers``, ``dev``). Note dev mode's 1000-patient cap is + applied downstream of ``load_table`` (in + ``BaseDataset._event_transform``), so it composes with + ``subset`` with no extra handling here. + + Examples: + >>> from pyhealth.datasets import MEDSDataset + >>> dataset = MEDSDataset( + ... root="/path/to/mimic-iv-demo-meds/0.0.1", + ... ) # doctest: +SKIP + >>> dataset.stats() # doctest: +SKIP + >>> # Canonical training split only, split map exposed as events: + >>> train = MEDSDataset( + ... root="/path/to/mimic-iv-demo-meds/0.0.1", + ... tables=["meds", "subject_splits"], + ... subset="train", + ... ) # doctest: +SKIP + """ + + def __init__( + self, + root: str, + tables: Optional[List[str]] = None, + subset: str = "all", + split_source: SplitSource = "metadata", + dataset_name: Optional[str] = None, + config_path: Optional[str] = None, + **kwargs, + ) -> None: + if subset not in (*MEDS_SPLITS, "all"): + raise ValueError( + f"subset must be one of {(*MEDS_SPLITS, 'all')}, got {subset!r}" + ) + if split_source not in ("metadata", "directory"): + raise ValueError( + f"split_source must be 'metadata' or 'directory', got {split_source!r}" + ) + + # Set before super().__init__: _init_cache_dir (called by the base + # constructor) reads them. + self.subset = subset + self.split_source = split_source + self._subset_patient_ids_cache: Optional[List[str]] = None + + if config_path is None: + logger.info("No config path provided, using default MEDS config") + config_path = Path(__file__).parent / "configs" / "meds.yaml" + + if tables is None: + tables = ["meds"] + + super().__init__( + root=root, + tables=tables, + dataset_name=dataset_name or "meds", + config_path=config_path, + **kwargs, + ) + + # Fail fast on schema-contract violations (footer read only). + self._validate_event_schema() + + # ------------------------------------------------------------------ + # Cache keying + # ------------------------------------------------------------------ + + def _init_cache_dir(self, cache_dir) -> Path: + """Nest a subset-specific directory under the standard cache key. + + The base cache key hashes only ``{root, tables, dataset_name, dev}`` + (``BaseDataset._init_cache_dir``); ``subset`` changes the *content* + of the cached ``global_event_df`` because rows are filtered in + ``load_data``, so instances with different subsets (or different + split sources) must not share a cache. ``subset="all"`` (default) + keeps the exact upstream cache layout. + """ + base = super()._init_cache_dir(cache_dir) + if self.subset == "all": + return base + sub = base / f"subset-{self.split_source}-{self.subset}" + sub.mkdir(parents=True, exist_ok=True) + return sub + + # ------------------------------------------------------------------ + # Split handling (canonical split as events + optional subset filter) + # ------------------------------------------------------------------ + + def _subset_patient_ids(self) -> Optional[List[str]]: + """Patient IDs belonging to ``self.subset``; ``None`` for ``"all"``. + + * ``split_source="metadata"``: read the canonical mapping file. One + row per subject, so plain pandas is enough. Column names + ``subject_id`` / ``split``. + * ``split_source="directory"``: subjects found under + ``data//``. Column projection keeps the read cheap. + + Computed once per instance and reused across tables, so multi-table + loads pay the read a single time. + """ + if self.subset == "all": + return None + if self._subset_patient_ids_cache is None: + if self.split_source == "metadata": + path = Path(clean_path(f"{self.root}/{SUBJECT_SPLITS_RELPATH}")) + if not path.exists(): + raise FileNotFoundError( + f"subset={self.subset!r} with split_source='metadata' " + f"requires {SUBJECT_SPLITS_RELPATH} under " + f"{self.root!r}. Pass split_source='directory' to " + "derive the split from the data// layout, or " + "use subset='all'." + ) + splits = pd.read_parquet(path).rename(columns=str.lower) + ids = splits.loc[splits["split"] == self.subset, "subject_id"] + else: # "directory" + split_dir = Path( + clean_path(f"{self.root}/{DATA_RELPATH}/{self.subset}") + ) + if not split_dir.is_dir(): + raise FileNotFoundError( + f"subset={self.subset!r} with split_source=" + f"'directory' requires the directory " + f"{DATA_RELPATH}/{self.subset} under {self.root!r}." + ) + ids = ( + self._scan_parquet(str(split_dir))["subject_id"].unique().compute() + ) + self._subset_patient_ids_cache = ids.astype("string").dropna().tolist() + logger.info( + f"MEDS subset={self.subset!r} via split_source=" + f"{self.split_source!r}: " + f"{len(self._subset_patient_ids_cache)} patients" + ) + return self._subset_patient_ids_cache + + def load_data(self) -> dd.DataFrame: + """Load all configured tables, restricted to the subset if any. + + Returns: + dd.DataFrame: The concatenated event frame, filtered to the + subjects of ``self.subset`` when a split was requested. + """ + df = super().load_data() + subset_ids = self._subset_patient_ids() + if subset_ids is not None: + df = df[df["patient_id"].isin(subset_ids)] + return df + + def _validate_event_schema(self) -> None: + """Fails fast when a Parquet event table violates the MEDS contract. + + Only Parquet footers are read (no data, no Dask). For every selected + table whose source is Parquet and whose timestamp is a single column, + that column must exist and be a timezone-naive timestamp type: the + MEDS reference ``DataSchema`` defines ``time`` as ``timestamp[us]`` + without a timezone (verified against the ``meds`` 0.4.1 package). + + This closes, at construction time, the silent-parse hazard of + date-like integers: an ``int64`` column holding ``20240101`` is + rejected here by dtype instead of being parsed as a date deep + inside the Dask graph. + + Raises: + TypeError: If the timestamp column is missing from the Parquet + schema, is not a timestamp type, or is timezone-aware. + """ + for name in self.tables: + table_cfg = self.config.tables.get(name.lower()) + if table_cfg is None: + continue # unknown table: load_table raises the proper error + ts_col = table_cfg.timestamp + if not ts_col or isinstance(ts_col, list): + continue + source = Path(clean_path(f"{self.root}/{table_cfg.file_path}")) + if source.suffix not in (".parquet", ".pq") and not source.is_dir(): + continue # non-Parquet source: string-parse contract applies + schema = pa_ds.dataset(str(source), format="parquet").schema + fields = {field.name.lower(): field for field in schema} + field = fields.get(ts_col.lower()) + if field is None: + raise TypeError( + f"MEDS table '{name}': timestamp column '{ts_col}' is " + f"missing from the Parquet schema {schema.names}." + ) + if not pa.types.is_timestamp(field.type): + raise TypeError( + f"MEDS table '{name}': column '{ts_col}' must be a " + f"timestamp in the Parquet schema, got '{field.type}'. " + "Date-like integers or strings parse unreliably; convert " + "the column upstream (e.g. with MEDS-Transform)." + ) + if field.type.tz is not None: + raise TypeError( + f"MEDS table '{name}': column '{ts_col}' is timezone-" + f"aware ('{field.type}'), but the MEDS reference schema " + "is timezone-naive (timestamp[us]). Normalize upstream, " + "e.g. tz_convert('UTC').tz_localize(None)." + ) diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 406b457f2..df8411db0 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -22,6 +22,9 @@ drug_recommendation_mimic4_fn, drug_recommendation_omop_fn, ) +from .in_hospital_mortality_meds import ( + InHospitalMortalityMEDS as InHospitalMortalityMEDS, +) from .in_hospital_mortality_mimic4 import InHospitalMortalityMIMIC4 from .length_of_stay_prediction import ( LengthOfStayPredictioneICU, diff --git a/pyhealth/tasks/in_hospital_mortality_meds.py b/pyhealth/tasks/in_hospital_mortality_meds.py new file mode 100644 index 000000000..7c9a9247d --- /dev/null +++ b/pyhealth/tasks/in_hospital_mortality_meds.py @@ -0,0 +1,252 @@ +"""In-hospital mortality prediction for datasets in the Medical Event Data +Standard (MEDS). + +This module provides :class:`InHospitalMortalityMEDS`, the MEDS-native +counterpart of +:class:`~pyhealth.tasks.in_hospital_mortality_mimic4.InHospitalMortalityMIMIC4`. +The MIMIC-IV task anchors on a visit object and reads +``admission.hospital_expire_flag``; MEDS represents a hospitalization as two +separate events (``HOSPITAL_ADMISSION//*`` and ``HOSPITAL_DISCHARGE//*``) +that share a ``hadm_id``, so a stay is derived by grouping those events +on that identifier and the label is derived from the discharge code. + +Task definition +--------------- +Let a *stay* be the set of events sharing one ``hadm_id`` for a subject, +with admission time ``t_a`` (earliest admission event) and discharge time +``t_d`` (latest discharge event). For each completed stay +(``t_d > t_a``) the task produces one sample: + +* **Prediction time** ``t_p``. + ``observation_window="full_stay"`` (default) sets ``t_p = t_d``. + ``observation_window="first_hours"`` sets ``t_p = t_a + window_hours`` and + keeps only stays with length of stay strictly greater than + ``window_hours``, so the window is fully observed and the outcome is + strictly future. +* **Features.** The ordered sequence of MEDS ``code`` values in the + half-open interval ``[t_a, t_p)``, excluding every ``HOSPITAL_DISCHARGE//*`` + event and every ``MEDS_DEATH`` event. Both exclusions matter: the half-open + bound already removes the discharge event when ``t_p = t_d``, and dropping + ``MEDS_DEATH`` removes the canonical death sentinel (which the demo places a + few hours after discharge) so the outcome can never leak into the input. +* **Label.** ``mortality = 1`` iff the stay's discharge code is + ``HOSPITAL_DISCHARGE//DIED``. This is the in-hospital, same-stay + definition, consistent with ``hospital_expire_flag`` upstream. On the + public MIMIC-IV demo in MEDS it is a strict superset of ``MEDS_DEATH`` + occurring within a stay: ``MEDS_DEATH`` carries a null ``hadm_id`` there, + so it cannot be attached to a stay and is deliberately not used as the + label. Deaths outside the index stay are a subject-level problem and are + out of scope for this task. + +Configuration +------------- +The task reads ``hadm_id``, which is **not** part of the core MEDS schema +(``subject_id``/``time``/``code``/``numeric_value``/``text_value``) but is +present in MIMIC-derived MEDS datasets. It is therefore kept out of the +default ``configs/meds.yaml`` (selecting an absent attribute would raise for +generic MEDS data). A bundled ``configs/meds_with_hadm.yaml`` exposes it; +pass that config (or your own that lists ``hadm_id``) when using this task. + +Scope note +---------- +The MIMIC-IV task additionally drops pediatric admissions via +``anchor_age``. A MEDS-native age filter is derivable from ``MEDS_BIRTH`` but +is intentionally omitted here: its on-disk representation is not fixed across +MEDS datasets, and silently assuming one would be unsound. Age restriction is +therefore left to a preprocessing step or a future, explicitly parameterized +extension. + +References: + MEDS Working Group. Medical Event Data Standard (MEDS): Facilitating + Machine Learning for Health. ICLR 2024 Workshop on Learning from Time + Series For Health. https://openreview.net/forum?id=IsHy2ebjIG +""" + +from typing import Any, ClassVar, Dict, List, Optional, Tuple + +import polars as pl + +from .base_task import BaseTask + +ADMISSION_PREFIX = "HOSPITAL_ADMISSION" +DISCHARGE_PREFIX = "HOSPITAL_DISCHARGE" +DIED_CODE = "HOSPITAL_DISCHARGE//DIED" +DEATH_CODE = "MEDS_DEATH" + +_FULL_STAY = "full_stay" +_FIRST_HOURS = "first_hours" +_VALID_WINDOWS = (_FULL_STAY, _FIRST_HOURS) + + +class InHospitalMortalityMEDS(BaseTask): + """In-hospital mortality prediction for MEDS datasets. + + One sample per completed hospital stay. The observation window is the + half-open interval ``[admission, prediction_time)`` and the binary label + is whether the stay ended in death (discharge code + ``HOSPITAL_DISCHARGE//DIED``). MEDS codes observed during the window, + excluding the terminating discharge event and any ``MEDS_DEATH``, form + the input sequence. See the module docstring for the full definition. + + Args: + observation_window (str): ``"full_stay"`` (default) observes the + entire stay, i.e. ``[admission, discharge)``. ``"first_hours"`` + observes only ``[admission, admission + window_hours)`` and keeps + stays whose length exceeds ``window_hours`` (an early-warning + setup with a strictly future outcome). + window_hours (float): Observation length used when + ``observation_window="first_hours"``. Ignored for ``"full_stay"``. + Defaults to ``48.0``, matching ``InHospitalMortalityMIMIC4``. + code_mapping (Optional[Dict[str, Tuple[str, str]]]): Optional vocab + mapping forwarded to :class:`BaseTask` (e.g. + ``{"codes": ("ICD10CM", "CCSCM")}``). + + Attributes: + task_name (str): The name of the task. + input_schema (Dict[str, str]): ``codes`` — the sequence of MEDS + codes observed during the window. + output_schema (Dict[str, str]): ``mortality`` — binary in-hospital + mortality. + + Raises: + ValueError: If ``observation_window`` is not one of + ``"full_stay"``/``"first_hours"``, or if ``window_hours`` is not + positive. + + Examples: + >>> from pathlib import Path + >>> import pyhealth.datasets.configs as meds_configs + >>> from pyhealth.datasets import MEDSDataset + >>> from pyhealth.tasks import InHospitalMortalityMEDS + >>> # A bundled stay-aware config exposes hadm_id (not a core MEDS + >>> # field, so it is kept out of the default configs/meds.yaml): + >>> cfg = Path(meds_configs.__file__).parent / "meds_with_hadm.yaml" + >>> dataset = MEDSDataset( + ... root="/path/to/mimic-iv-demo-meds/0.0.1", + ... config_path=str(cfg), + ... ) + >>> samples = dataset.set_task(InHospitalMortalityMEDS()) + >>> # Early-warning variant: first 48h, stays longer than 48h only + >>> early = InHospitalMortalityMEDS(observation_window="first_hours") + """ + + task_name: str = "InHospitalMortalityMEDS" + input_schema: ClassVar[Dict[str, str]] = {"codes": "sequence"} + output_schema: ClassVar[Dict[str, str]] = {"mortality": "binary"} + + def __init__( + self, + observation_window: str = _FULL_STAY, + window_hours: float = 48.0, + code_mapping: Optional[Dict[str, Tuple[str, str]]] = None, + ) -> None: + if observation_window not in _VALID_WINDOWS: + raise ValueError( + f"observation_window must be one of {_VALID_WINDOWS}, " + f"got {observation_window!r}." + ) + if window_hours <= 0: + raise ValueError(f"window_hours must be positive, got {window_hours}.") + super().__init__(code_mapping=code_mapping) + self.observation_window = observation_window + self.window_hours = float(window_hours) + + def pre_filter(self, df: pl.LazyFrame) -> pl.LazyFrame: + """Restricts the global scan to MEDS events before per-patient calls. + + All MEDS data lives in a single ``meds`` event type, so this narrows + the frame once rather than per patient. + """ + return df.filter(pl.col("event_type") == "meds") + + def _group_stays(self, events: pl.DataFrame) -> pl.DataFrame: + """Builds one row per stay from admission/discharge events. + + Args: + events (pl.DataFrame): This patient's MEDS events, with an + integer ``_hadm`` column already attached. + + Returns: + pl.DataFrame: Columns ``_hadm``, ``admit``, ``discharge``, + ``discharge_code``, one row per ``hadm_id`` that has both an + admission and a discharge. Malformed duplicates collapse via + earliest-admission / latest-discharge aggregation. + """ + code = pl.col("meds/code") + admissions = ( + events.filter(code.str.starts_with(ADMISSION_PREFIX)) + .filter(pl.col("_hadm").is_not_null()) + .group_by("_hadm") + .agg(pl.col("timestamp").min().alias("admit")) + ) + discharges = ( + events.filter(code.str.starts_with(DISCHARGE_PREFIX)) + .filter(pl.col("_hadm").is_not_null()) + .group_by("_hadm") + .agg( + pl.col("timestamp").max().alias("discharge"), + code.sort_by("timestamp").last().alias("discharge_code"), + ) + ) + return admissions.join(discharges, on="_hadm", how="inner") + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + events = patient.get_events(event_type="meds", return_df=True) + if events.height == 0: + return [] + + # A nullable integer id is promoted to float through the Dask/pandas + # pipeline whenever the column carries nulls (e.g. lab events and the + # MEDS_DEATH sentinel). Cast back to a nullable integer so stays join + # cleanly and emitted ids stay integral rather than "555.0". + events = events.with_columns( + pl.col("meds/hadm_id").cast(pl.Int64, strict=False).alias("_hadm") + ) + code = pl.col("meds/code") + + stays = self._group_stays(events) + if stays.height == 0: + return [] + + samples: List[Dict[str, Any]] = [] + for stay in stays.sort("admit").iter_rows(named=True): + admit, discharge = stay["admit"], stay["discharge"] + if discharge <= admit: + continue # degenerate/zero-length stay + + if self.observation_window == _FIRST_HOURS: + duration_hours = (discharge - admit).total_seconds() / 3600.0 + if duration_hours <= self.window_hours: + continue # window not fully observed within this stay + predict_time = admit + _timedelta_hours(self.window_hours) + else: + predict_time = discharge + + window = events.filter( + (pl.col("timestamp") >= admit) + & (pl.col("timestamp") < predict_time) # half-open: excludes t_p + & (~code.str.starts_with(DISCHARGE_PREFIX)) + & (code != DEATH_CODE) + ).sort("timestamp") + + codes = window["meds/code"].to_list() + if not codes: + continue # no observable signal before the prediction time + + samples.append( + { + "patient_id": patient.patient_id, + "hadm_id": stay["_hadm"], + "codes": codes, + "mortality": int(stay["discharge_code"] == DIED_CODE), + } + ) + + return samples + + +def _timedelta_hours(hours: float): + """Returns a ``datetime.timedelta`` of ``hours`` (kept import-local).""" + from datetime import timedelta + + return timedelta(hours=hours) diff --git a/tests/core/test_in_hospital_mortality_meds.py b/tests/core/test_in_hospital_mortality_meds.py new file mode 100644 index 000000000..6ae62aae7 --- /dev/null +++ b/tests/core/test_in_hospital_mortality_meds.py @@ -0,0 +1,278 @@ +"""Tests for InHospitalMortalityMEDS. + +The synthetic tests build a small MEDS-shaped dataset (typed Parquet shards +with a ``hadm_id`` column, plus a stay-aware config) with no real data and no +downloads. Feature-content assertions apply the task directly to real +``Patient`` objects (``dataset.get_patient(...)``), because that yields the +task's raw ``List[Dict]`` output; going through ``set_task`` instead would +tokenize the ``codes`` sequence into integer indices (the ``sequence`` +processor), which is the correct user path but hides the string codes the +leakage tests must inspect. A separate test exercises ``set_task`` end to end +to confirm the intended path produces the expected number of samples. + +The central property under test is the absence of label leakage: for a stay +that ends in death, neither the discharge event nor the ``MEDS_DEATH`` +sentinel may appear in the emitted feature sequence. +``TestInHospitalMortalityMEDSDemoSmoke`` optionally exercises the public +MIMIC-IV demo in MEDS when it is available locally, and is skipped otherwise. +""" + +import os +import shutil +import tempfile +import unittest +from datetime import datetime, timedelta +from pathlib import Path + +import polars as pl + +from pyhealth.datasets import MEDSDataset +from pyhealth.tasks import InHospitalMortalityMEDS +from pyhealth.tasks.in_hospital_mortality_meds import ( + DEATH_CODE, + DISCHARGE_PREFIX, +) + +T0 = datetime(2024, 1, 1, 8, 0, 0) + +# Stay-aware config: exposes hadm_id, which the task requires. +_CONFIG = """version: "1.0" +tables: + meds: + file_path: "data" + patient_id: "subject_id" + timestamp: "time" + attributes: + - "code" + - "numeric_value" + - "hadm_id" +""" + + +def _events_to_frame(rows): + """Rows are (subject_id, offset_hours, code, hadm_id_or_None).""" + return pl.DataFrame( + { + "subject_id": pl.Series([r[0] for r in rows], dtype=pl.Int64), + "time": pl.Series( + [T0 + timedelta(hours=r[1]) for r in rows], dtype=pl.Datetime("us") + ), + "code": pl.Series([r[2] for r in rows], dtype=pl.String), + "numeric_value": pl.Series([None] * len(rows), dtype=pl.Float32), + "hadm_id": pl.Series( + [r[3] for r in rows], + dtype=pl.Int64, # nullable + ), + } + ) + + +class TestInHospitalMortalityMEDS(unittest.TestCase): + """Task behavior on a synthetic MEDS dataset via set_task.""" + + def setUp(self): + self.temp_dir = Path(tempfile.mkdtemp()) + self.root = self.temp_dir / "meds" + (self.root / "data").mkdir(parents=True) + self.cache_root = self.temp_dir / "cache" + self.config_path = self.temp_dir / "meds_hadm.yaml" + self.config_path.write_text(_CONFIG) + self._write_default_cohort() + + def tearDown(self): + if self.temp_dir.exists(): + # ignore_errors: litdata may keep chunk handles open on Windows. + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _write_default_cohort(self): + # Subject 1: a stay ending in death (hadm 555), a MEDS_DEATH a few + # hours later (null hadm), a post-death stray event, then a second + # stay ending at home (hadm 777). Subject 2: a survived stay whose + # subject later dies out of hospital. + rows = [ + (1, 0, "HOSPITAL_ADMISSION//EW", 555), + (1, 2, "LAB//50912", 555), + (1, 6, "MED//aspirin", 555), + (1, 10, "HOSPITAL_DISCHARGE//DIED", 555), + (1, 14, DEATH_CODE, None), + (1, 16, "LAB//stray", None), + (1, 120, "HOSPITAL_ADMISSION//OBS", 777), + (1, 121, "LAB//x", 777), + (1, 144, "HOSPITAL_DISCHARGE//HOME", 777), + (2, 0, "HOSPITAL_ADMISSION//EW", 999), + (2, 5, "HOSPITAL_DISCHARGE//HOME", 999), + (2, 200, DEATH_CODE, None), + ] + _events_to_frame(rows).write_parquet(self.root / "data" / "0.parquet") + + def _dataset(self): + return MEDSDataset( + root=str(self.root), + config_path=str(self.config_path), + cache_dir=self.cache_root, + ) + + def _apply(self, task=None): + """Applies the task to every patient, returning raw sample dicts. + + This mirrors what ``set_task`` does per patient but keeps the task's + untokenized output so feature sequences remain inspectable. + """ + task = task or InHospitalMortalityMEDS() + dataset = self._dataset() + samples = [] + for pid in dataset.unique_patient_ids: + samples.extend(task(dataset.get_patient(pid))) + return samples + + def _by_hadm(self, samples): + return {s["hadm_id"]: s for s in samples} + + def test_one_sample_per_completed_stay(self): + samples = self._apply() + # Three completed stays across the two subjects. + self.assertEqual(len(samples), 3) + self.assertEqual(sorted(s["hadm_id"] for s in samples), [555, 777, 999]) + # ids are integral, not promoted floats + self.assertTrue(all(isinstance(s["hadm_id"], int) for s in samples)) + + def test_labels_from_discharge_code(self): + by = self._by_hadm(self._apply()) + self.assertEqual(by[555]["mortality"], 1) # DIED + self.assertEqual(by[777]["mortality"], 0) # HOME + self.assertEqual(by[999]["mortality"], 0) # HOME (subject dies later) + + def test_no_label_leakage_in_positive_stay(self): + """The defining safety property: outcome never enters the features.""" + died = self._by_hadm(self._apply())[555] + for code in died["codes"]: + self.assertFalse(code.startswith(DISCHARGE_PREFIX)) + self.assertNotEqual(code, DEATH_CODE) + # Exactly the pre-discharge, non-death events, in order. + self.assertEqual( + died["codes"], + ["HOSPITAL_ADMISSION//EW", "LAB//50912", "MED//aspirin"], + ) + # The stray post-death event is excluded too. + self.assertNotIn("LAB//stray", died["codes"]) + + def test_meds_death_without_hadm_never_labels_a_stay(self): + # Subject 2 dies out of hospital; the in-hospital stay stays negative. + self.assertEqual(self._by_hadm(self._apply())[999]["mortality"], 0) + + def test_first_hours_requires_sufficient_length_of_stay(self): + # Default cohort: no stay exceeds 48h, so the early-warning variant + # yields nothing. + task = InHospitalMortalityMEDS(observation_window="first_hours") + self.assertEqual(self._apply(task), []) + + def test_first_hours_observes_only_the_window(self): + # A single long stay (LOS 100h) observed for its first 48h. + (self.root / "data" / "0.parquet").unlink() + rows = [ + (7, 0, "HOSPITAL_ADMISSION//EW", 900), + (7, 12, "LAB//a", 900), + (7, 47, "LAB//b", 900), # inside 48h + (7, 60, "LAB//c", 900), # outside 48h + (7, 100, "HOSPITAL_DISCHARGE//DIED", 900), + ] + _events_to_frame(rows).write_parquet(self.root / "data" / "0.parquet") + task = InHospitalMortalityMEDS( + observation_window="first_hours", window_hours=48.0 + ) + samples = self._apply(task) + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["mortality"], 1) # eventual outcome + self.assertEqual( + samples[0]["codes"], + ["HOSPITAL_ADMISSION//EW", "LAB//a", "LAB//b"], + ) + self.assertNotIn("LAB//c", samples[0]["codes"]) + + def test_discharge_boundary_is_half_open(self): + (self.root / "data" / "0.parquet").unlink() + rows = [ + (8, 0, "HOSPITAL_ADMISSION//EW", 111), + (8, 4, "LAB//inside", 111), + (8, 5, "LAB//at_discharge", 111), # exactly at t_discharge + (8, 5, "HOSPITAL_DISCHARGE//HOME", 111), + ] + _events_to_frame(rows).write_parquet(self.root / "data" / "0.parquet") + samples = self._apply() + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["codes"], ["HOSPITAL_ADMISSION//EW", "LAB//inside"]) + + def test_invalid_parameters_raise(self): + with self.assertRaises(ValueError): + InHospitalMortalityMEDS(observation_window="bogus") + with self.assertRaises(ValueError): + InHospitalMortalityMEDS(window_hours=0) + with self.assertRaises(ValueError): + InHospitalMortalityMEDS(window_hours=-3) + + def test_set_task_integration_yields_expected_count(self): + """The intended user path runs and produces one sample per stay. + + Codes are tokenized by the sequence processor here, so only the + sample count (structure) is asserted; content/leakage is covered by + the get_patient-based tests above. + """ + dataset = self._dataset() + sample_dataset = dataset.set_task(InHospitalMortalityMEDS()) + self.assertEqual(len(sample_dataset), 3) + + +def _demo_root() -> str: + env = os.environ.get("MEDS_DEMO_ROOT") + if env: + return env + test_dir = Path(__file__).parent.parent.parent + return str(test_dir / "test-resources" / "meds_demo") + + +@unittest.skipUnless( + Path(_demo_root()).is_dir(), + "MIMIC-IV demo in MEDS format not available locally " + "(set MEDS_DEMO_ROOT or place it under test-resources/meds_demo)", +) +class TestInHospitalMortalityMEDSDemoSmoke(unittest.TestCase): + """Smoke test on the public MIMIC-IV demo in MEDS format. + + Requires a config exposing hadm_id; this test writes one next to a + temporary cache. The demo (PhysioNet, https://doi.org/10.13026/t2y8-ea41, + ODbL v1.0) is never downloaded here. + """ + + def setUp(self): + self.temp_dir = Path(tempfile.mkdtemp()) + self.config_path = self.temp_dir / "meds_hadm.yaml" + self.config_path.write_text(_CONFIG) + + def tearDown(self): + if self.temp_dir.exists(): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_produces_stays_without_leakage(self): + dataset = MEDSDataset( + root=_demo_root(), + config_path=str(self.config_path), + cache_dir=self.temp_dir, + ) + task = InHospitalMortalityMEDS() + # Apply per patient to inspect raw (untokenized) code sequences. + samples = [] + for pid in dataset.unique_patient_ids: + samples.extend(task(dataset.get_patient(pid))) + self.assertGreater(len(samples), 0) + # No sample may contain a discharge or death code (leakage guard). + for sample in samples: + for code in sample["codes"]: + self.assertFalse(code.startswith(DISCHARGE_PREFIX)) + self.assertNotEqual(code, DEATH_CODE) + n_positive = sum(int(s["mortality"]) for s in samples) + self.assertGreater(n_positive, 0) + self.assertLess(n_positive, len(samples)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_meds.py b/tests/core/test_meds.py new file mode 100644 index 000000000..4ad39658c --- /dev/null +++ b/tests/core/test_meds.py @@ -0,0 +1,309 @@ +"""Tests for MEDSDataset (synthetic Parquet fixtures, fixed seeds). + +Optional smoke on a real export: set ``MEDS_DEMO_ROOT`` to the dataset version +directory (the folder containing ``data/`` and ``metadata/``). +""" + +from __future__ import annotations + +import os +import shutil +import tempfile +import unittest +from datetime import datetime +from pathlib import Path +from typing import Dict, List +from unittest.mock import patch + +import numpy as np +import pandas as pd +import polars as pl +import pyarrow as pa +import pyarrow.parquet as pq + +from pyhealth.datasets import MEDSDataset, MIMIC3Dataset +from pyhealth.tasks.base_task import BaseTask + +T1 = datetime(2024, 1, 1, 8, 0, 0) +T2 = datetime(2024, 1, 2, 9, 30, 0) + +MEDS_DEMO_ROOT = os.environ.get("MEDS_DEMO_ROOT") +_DEFAULT_DEMO = ( + Path(__file__).parent.parent.parent / "test-resources" / "meds_demo" +) +MEDS_DEMO_PATH = ( + Path(MEDS_DEMO_ROOT).expanduser() + if MEDS_DEMO_ROOT + else _DEFAULT_DEMO +) + +# Fixed split assignment for deterministic assertions. +_SYNTHETIC_SPLITS: Dict[str, List[int]] = { + "train": [1001, 1002, 1003, 1004], + "tuning": [1005, 1006], + "held_out": [1007, 1008], +} + + +def write_synthetic_meds(root: Path, *, seed: int = 42, rows_per_shard: int = 25) -> None: + """Write a MEDS-shaped tree under ``root`` (``data//*.parquet`` + metadata).""" + rng = np.random.default_rng(seed) + data_root = root / "data" + for split, subjects in _SYNTHETIC_SPLITS.items(): + split_dir = data_root / split + split_dir.mkdir(parents=True, exist_ok=True) + for shard in range(2): + n = rows_per_shard + pq.write_table( + pa.table( + { + "subject_id": pa.array(rng.choice(subjects, n), type=pa.int64()), + "time": pa.array( + pd.date_range("2020-01-01", periods=n, freq="h"), + type=pa.timestamp("us"), + ), + "code": pa.array( + [f"LAB//{i % 5}" for i in range(n)], type=pa.string() + ), + "numeric_value": pa.array( + rng.normal(size=n).astype(np.float32), type=pa.float32() + ), + } + ), + split_dir / f"{shard}.parquet", + ) + + meta = root / "metadata" + meta.mkdir(exist_ok=True) + all_subjects = [sid for ids in _SYNTHETIC_SPLITS.values() for sid in ids] + all_splits = [ + split for split, ids in _SYNTHETIC_SPLITS.items() for _ in ids + ] + pq.write_table( + pa.table( + { + "subject_id": pa.array(all_subjects, type=pa.int64()), + "split": pa.array(all_splits, type=pa.string()), + } + ), + meta / "subject_splits.parquet", + ) + + +class _MedsSmokeTask(BaseTask): + """Minimal task: one sample per patient with at least one meds event.""" + + task_name: str = "MedsSmokeTask" + input_schema: Dict[str, str] = {"codes": "sequence"} + output_schema: Dict[str, str] = {"has_events": "binary"} + + def __call__(self, patient): + meds = patient.get_events(event_type="meds") + if not meds: + return [] + codes = [event.code for event in meds if event.code] + if not codes: + return [] + return [ + { + "patient_id": patient.patient_id, + "codes": codes, + "has_events": int(patient.patient_id) % 2, + } + ] + + +class TestMEDSDatasetSynthetic(unittest.TestCase): + """MEDSDataset against a local synthetic MEDS export.""" + + @classmethod + def setUpClass(cls) -> None: + cls._tmp = tempfile.mkdtemp(prefix="meds_synthetic_") + cls.root = Path(cls._tmp) + write_synthetic_meds(cls.root) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls._tmp, ignore_errors=True) + + def _dataset(self, **kwargs) -> MEDSDataset: + return MEDSDataset( + root=str(self.root), + cache_dir=self._tmp, + num_workers=1, + **kwargs, + ) + + def test_load_table_schema_and_dtypes(self) -> None: + ds = self._dataset() + df = ds.load_table("meds").compute() + self.assertIn("patient_id", df.columns) + self.assertIn("timestamp", df.columns) + self.assertIn("event_type", df.columns) + self.assertIn("meds/code", df.columns) + self.assertIn("meds/numeric_value", df.columns) + self.assertEqual(str(df["patient_id"].dtype), "string") + self.assertEqual(str(df["timestamp"].dtype), "datetime64[ms]") + self.assertTrue((df["event_type"] == "meds").all()) + + def test_loads_all_patients(self) -> None: + ds = self._dataset() + expected = {str(sid) for sid in _SYNTHETIC_SPLITS["train"]} + expected |= {str(sid) for sid in _SYNTHETIC_SPLITS["tuning"]} + expected |= {str(sid) for sid in _SYNTHETIC_SPLITS["held_out"]} + self.assertEqual(set(ds.unique_patient_ids), expected) + + def test_subset_train_via_metadata(self) -> None: + ds = self._dataset(subset="train", split_source="metadata") + expected = {str(sid) for sid in _SYNTHETIC_SPLITS["train"]} + self.assertEqual(set(ds.unique_patient_ids), expected) + + def test_subset_tuning_via_directory(self) -> None: + ds = self._dataset(subset="tuning", split_source="directory") + expected = {str(sid) for sid in _SYNTHETIC_SPLITS["tuning"]} + self.assertEqual(set(ds.unique_patient_ids), expected) + + def test_subject_splits_exposed_as_events(self) -> None: + ds = self._dataset(tables=["meds", "subject_splits"]) + patient_id = str(_SYNTHETIC_SPLITS["train"][0]) + patient = ds.get_patient(patient_id) + split_events = patient.get_events(event_type="subject_splits") + self.assertEqual(len(split_events), 1) + self.assertEqual(split_events[0].split, "train") + + def test_patient_meds_event_attributes(self) -> None: + ds = self._dataset(subset="train") + patient = ds.get_patient(str(_SYNTHETIC_SPLITS["train"][0])) + meds = patient.get_events(event_type="meds") + self.assertGreater(len(meds), 0) + self.assertTrue(str(meds[0].code).startswith("LAB//")) + + def test_invalid_subset_raises(self) -> None: + with self.assertRaises(ValueError): + self._dataset(subset="validation") + + def test_schema_violations_raise_type_error_at_construction(self): + """The footer guard rejects non-conforming `time` before any Dask. + + Covers the ADR 002 T5 hazard (date-like ints such as 20240101 would + otherwise parse silently) plus strings, timezone-aware timestamps, + and a missing column. + """ + cases = { + "int64": ( + "time", + pl.Series([20240101, 20240102], dtype=pl.Int64), + ), + "string": ( + "time", + pl.Series(["2024-01-01", "2024-01-02"], dtype=pl.String), + ), + "tz_aware": ( + "time", + pl.Series([T1, T2], dtype=pl.Datetime("us", "UTC")), + ), + "missing": ("ts", pl.Series([T1, T2], dtype=pl.Datetime("us"))), + } + for label, (col_name, series) in cases.items(): + with self.subTest(time=label): + bad_root = Path(self._tmp) / f"meds_bad_{label}" + (bad_root / "data").mkdir(parents=True) + pl.DataFrame( + { + "subject_id": pl.Series([1, 2], dtype=pl.Int64), + col_name: series, + "code": pl.Series(["A", "B"], dtype=pl.String), + "numeric_value": pl.Series( + [None, None], dtype=pl.Float32 + ), + } + ).write_parquet(bad_root / "data" / "0.parquet") + with self.assertRaises(TypeError): + MEDSDataset(root=str(bad_root), cache_dir=self._tmp) + + def test_cache_dir_varies_with_subset(self) -> None: + with patch( + "pyhealth.datasets.base_dataset.platformdirs.user_cache_dir", + return_value=self._tmp, + ): + all_ds = MEDSDataset( + root=str(self.root), + cache_dir=self._tmp, + num_workers=1, + ) + train_ds = MEDSDataset( + root=str(self.root), + cache_dir=self._tmp, + subset="train", + split_source="metadata", + num_workers=1, + ) + self.assertNotEqual(all_ds.cache_dir, train_ds.cache_dir) + + def test_set_task_smoke(self) -> None: + ds = self._dataset(subset="train") + sample_ds = ds.set_task(_MedsSmokeTask(), num_workers=1) + self.assertGreater(len(sample_ds), 0) + sample = sample_ds[0] + self.assertIn("codes", sample) + self.assertEqual(sample["has_events"], int(sample["patient_id"]) % 2) + + def test_mimic3_csv_path_unchanged(self) -> None: + """Non-regression: CSV-backed datasets still load after MEDSDataset addition.""" + demo = ( + Path(__file__).parent.parent.parent + / "test-resources" + / "core" + / "mimic3demo" + ) + ds = MIMIC3Dataset( + root=str(demo), + tables=["diagnoses_icd"], + cache_dir=self._tmp, + num_workers=1, + ) + self.assertGreater(len(ds.unique_patient_ids), 0) + + +@unittest.skipUnless( + MEDS_DEMO_PATH.is_dir() + and (MEDS_DEMO_PATH / "data").is_dir() + and (MEDS_DEMO_PATH / "metadata" / "subject_splits.parquet").is_file(), + "Download mimic-iv-demo-meds into test-resources/meds_demo or set MEDS_DEMO_ROOT", +) +class TestMEDSDatasetDemoSmoke(unittest.TestCase): + """Smoke on mimic-iv-demo-meds (partial export is enough for dtype checks).""" + + @classmethod + def setUpClass(cls) -> None: + cls.root = MEDS_DEMO_PATH.resolve() + cls.cache = tempfile.mkdtemp(prefix="meds_demo_") + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.cache, ignore_errors=True) + + def test_demo_load_table_dtypes(self) -> None: + ds = MEDSDataset( + root=str(self.root), + cache_dir=self.cache, + num_workers=1, + ) + df = ds.load_table("meds").compute() + self.assertEqual(str(df["patient_id"].dtype), "string") + self.assertEqual(str(df["timestamp"].dtype), "datetime64[ms]") + self.assertGreater(len(df), 0) + + def test_demo_stats_and_subset(self) -> None: + ds = MEDSDataset( + root=str(self.root), + cache_dir=self.cache, + subset="train", + num_workers=1, + ) + ds.stats() + self.assertGreater(len(ds.unique_patient_ids), 0) + + +if __name__ == "__main__": + unittest.main()