Add MEDS dataset support (MEDSDataset + typed Parquet scan path)

.gitignore

@@ -140,4 +140,18 @@ data/physionet.org/
 .codex
 
 # Model weight files (large binaries, distributed separately)
-weightfiles/
+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/

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

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) <https://github.com/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 <https://openreview.net/forum?id=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 <https://medical-event-data-standard.github.io/>`_.
+
+.. autoclass:: pyhealth.datasets.MEDSDataset
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/api/tasks.rst

@@ -207,6 +207,7 @@ Available Tasks
 
     Base Task <tasks/pyhealth.tasks.BaseTask>
     In-Hospital Mortality (MIMIC-IV) <tasks/pyhealth.tasks.InHospitalMortalityMIMIC4>
+    In-Hospital Mortality (MEDS) <tasks/pyhealth.tasks.InHospitalMortalityMEDS>
     MIMIC-III ICD-9 Coding <tasks/pyhealth.tasks.MIMIC3ICD9Coding>
     Cardiology Detection <tasks/pyhealth.tasks.cardiology_detect>
     COVID-19 CXR Classification <tasks/pyhealth.tasks.COVID19CXRClassification>

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:

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/<split>/<shard>.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()

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()

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