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..38d739939 --- /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 (Arnrich et al., ICLR 2024 Workshop on Learning from Time Series For Health). 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. + +.. 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..a34d5f774 --- /dev/null +++ b/examples/meds_demo.py @@ -0,0 +1,76 @@ +"""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 polars as pl + +from pyhealth.datasets import MEDSDataset + + +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}") + + # From here, the dataset behaves like any other PyHealth dataset: use + # `dataset.set_task(...)` with an existing task to build samples. + + +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..3efe5a90c --- /dev/null +++ b/examples/verify_meds_mortality.py @@ -0,0 +1,69 @@ +"""Standalone verification of the MEDS in-hospital mortality cohort. + +Prints the cohort summary produced by ``InHospitalMortalityMEDS`` on a MEDS +dataset, so the positive rate (expected ~12/238 on the public MIMIC-IV demo) +can be confirmed through the actual task pipeline rather than only at the raw +Parquet level. + +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 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, + ) + + # Apply per patient to keep raw (untokenized) samples, so the label + # counts are directly inspectable. set_task would tokenize the codes. + samples = [] + for patient_id in dataset.unique_patient_ids: + samples.extend(task(dataset.get_patient(patient_id))) + + summary = InHospitalMortalityMEDS.summarize(samples) + print(f"root : {args.root}") + print(f"observation_window : {args.observation_window}") + print(f"n_samples (stays) : {summary['n_samples']}") + print(f"n_patients : {summary['n_patients']}") + print(f"n_positive (died) : {summary['n_positive']}") + print(f"positive_rate : {summary['positive_rate']:.4f}") + print(f"mean_sequence_length : {summary['mean_sequence_length']:.1f}") + + # Cross-check the intended user path returns the same sample count. + sample_dataset = dataset.set_task(task) + print(f"set_task sample count : {len(sample_dataset)} (should equal n_samples)") + + +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..8eafe495d --- /dev/null +++ b/pyhealth/datasets/meds.py @@ -0,0 +1,301 @@ +"""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``. 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``. + + ``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 (see ADR). + + 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 (ADR 002, T5): 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..c6b115a24 --- /dev/null +++ b/pyhealth/tasks/in_hospital_mortality_meds.py @@ -0,0 +1,287 @@ +"""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 reconstructed by joining 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 _reconstruct_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._reconstruct_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 + + @staticmethod + def summarize(samples: List[Dict[str, Any]]) -> Dict[str, Any]: + """Summary statistics of a generated sample set, for cohort reporting. + + Args: + samples (List[Dict[str, Any]]): Output of applying this task. + + Returns: + Dict[str, Any]: ``n_samples``, ``n_patients``, ``n_positive``, + ``positive_rate`` (0.0 for an empty set), and + ``mean_sequence_length``. + + Examples: + >>> InHospitalMortalityMEDS.summarize([]) + {'n_samples': 0, 'n_patients': 0, 'n_positive': 0, \ +'positive_rate': 0.0, 'mean_sequence_length': 0.0} + """ + n = len(samples) + if n == 0: + return { + "n_samples": 0, + "n_patients": 0, + "n_positive": 0, + "positive_rate": 0.0, + "mean_sequence_length": 0.0, + } + n_positive = sum(int(s["mortality"]) for s in samples) + return { + "n_samples": n, + "n_patients": len({s["patient_id"] for s in samples}), + "n_positive": n_positive, + "positive_rate": n_positive / n, + "mean_sequence_length": sum(len(s["codes"]) for s in samples) / n, + } + + +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..3a1c8b040 --- /dev/null +++ b/tests/core/test_in_hospital_mortality_meds.py @@ -0,0 +1,290 @@ +"""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_summarize(self): + summ = InHospitalMortalityMEDS.summarize(self._apply()) + self.assertEqual(summ["n_samples"], 3) + self.assertEqual(summ["n_positive"], 1) + self.assertAlmostEqual(summ["positive_rate"], 1 / 3) + self.assertEqual(summ["n_patients"], 2) + empty = InHospitalMortalityMEDS.summarize([]) + self.assertEqual(empty["positive_rate"], 0.0) + + 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) + summ = InHospitalMortalityMEDS.summarize(samples) + self.assertGreater(summ["n_positive"], 0) + self.assertLess(summ["n_positive"], summ["n_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()