From e7c8bbbc2f04c5b4d5ffda4e2e5d4b34f6cfb6a9 Mon Sep 17 00:00:00 2001 From: solarsys Date: Fri, 2 Oct 2026 20:52:06 -0400 Subject: [PATCH 1/6] Fit processors for schema fields that were not supplied SampleBuilder fitted input (or output) processors only when none were supplied. Supplying a pre-fitted processor for one field left every other field on that side without a processor, so its raw value (e.g. a Python list) reached the model with no error. This is the common fit-on-train, reuse-for-test pattern when a user reuses only some processors. - SampleBuilder.fit fits a processor for every schema field that has none; supplied processors are used as they are, never refitted. - SampleBuilder copies the supplied processor dicts, so fitting never adds keys to the caller's dict (e.g. the training set's processors). - tests/core/test_partial_processors.py: missing input and output processors are fitted, supplied ones are not refitted, and create_sample_dataset gives tensors for every field. - docs/api/processors.rst: "Reusing Fitted Processors". - examples/reuse_train_processors.py: fit on train, reuse for test. Co-Authored-By: Claude Opus 5.5 --- docs/api/processors.rst | 25 ++++++++ examples/reuse_train_processors.py | 52 ++++++++++++++++ pyhealth/datasets/sample_dataset.py | 55 ++++++++++++----- tests/core/test_partial_processors.py | 85 +++++++++++++++++++++++++++ 4 files changed, 201 insertions(+), 16 deletions(-) create mode 100644 examples/reuse_train_processors.py create mode 100644 tests/core/test_partial_processors.py diff --git a/docs/api/processors.rst b/docs/api/processors.rst index e4262fa4c..783e34353 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -276,6 +276,31 @@ For maximum control, you can pass pre-configured processor instances directly in return samples +Reusing Fitted Processors +------------------------- + +``set_task()`` and ``create_sample_dataset()`` accept already-fitted processors +through ``input_processors`` and ``output_processors``. A supplied processor is +used as it is and never refitted; every schema field you do not supply gets a +processor fitted on the samples as usual. A common use is fitting on training +samples only and reusing those processors for validation and test data, so test +patients do not shape preprocessing: + +.. code-block:: python + + from pyhealth.datasets import create_sample_dataset + + train = create_sample_dataset(train_samples, input_schema, output_schema) + test = create_sample_dataset( + test_samples, input_schema, output_schema, + input_processors=train.input_processors, + output_processors=train.output_processors, + ) + +In PyHealth 2.0.2 and earlier, supplying a processor for some fields fitted +none for the rest of that schema: those fields reached the model as raw Python +values, without an error. See ``examples/reuse_train_processors.py``. + Processor String Keys --------------------- diff --git a/examples/reuse_train_processors.py b/examples/reuse_train_processors.py new file mode 100644 index 000000000..bdb8bafee --- /dev/null +++ b/examples/reuse_train_processors.py @@ -0,0 +1,52 @@ +"""Fit processors on training samples and reuse them for test samples. + +Fitting processors only on training data keeps test patients from shaping +preprocessing (vocabularies, and any statistics a processor learns). Pass the +training dataset's processors when building the test dataset: supplied +processors are reused as they are, never refitted. + +Runs in a few seconds on CPU with synthetic data; no download. + +Usage: + python examples/reuse_train_processors.py +""" + +from pyhealth.datasets import create_sample_dataset + +INPUT_SCHEMA = {"codes": "sequence", "age": "tensor"} +OUTPUT_SCHEMA = {"label": "binary"} + +train_samples = [ + {"patient_id": "p1", "codes": ["I10", "E11"], "age": [61.0], "label": 1}, + {"patient_id": "p2", "codes": ["J45"], "age": [34.0], "label": 0}, + {"patient_id": "p3", "codes": ["E11", "N18"], "age": [72.0], "label": 1}, +] +test_samples = [ + # "C50" never occurs in training, so it maps to , as at deployment + {"patient_id": "p9", "codes": ["I10", "C50"], "age": [58.0], "label": 0}, +] + + +def main(): + train = create_sample_dataset( + train_samples, INPUT_SCHEMA, OUTPUT_SCHEMA, dataset_name="train" + ) + codes = train.input_processors["codes"] + print("training vocabulary:", sorted(codes.code_vocab)) + + # Reuse every fitted processor. (A field left out of input_processors would + # get a processor fitted on the samples passed here, i.e. on test data.) + test = create_sample_dataset( + test_samples, + INPUT_SCHEMA, + OUTPUT_SCHEMA, + input_processors=train.input_processors, + output_processors=train.output_processors, + dataset_name="test", + ) + print("same codes processor:", test.input_processors["codes"] is codes) + print("test sample:", test[0]) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 6b9a751ee..6a6813003 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -49,6 +49,29 @@ class SampleBuilder: After saving the schema, `litdata.optimize` can be used with `builder.transform` to serialize and chunk pickled sample items into a directory that can be loaded via SampleDataset. + + Pre-fitted processors passed as `input_processors` / `output_processors` are + used as they are; every schema field without one gets a fitted processor. + + Examples: + >>> from pyhealth.datasets.sample_dataset import SampleBuilder + >>> from pyhealth.processors import SequenceProcessor + >>> samples = [ + ... {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1}, + ... {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0}, + ... ] + >>> codes = SequenceProcessor() + >>> codes.fit(samples, "codes") + >>> builder = SampleBuilder( + ... input_schema={"codes": "sequence", "age": "tensor"}, + ... output_schema={"label": "binary"}, + ... input_processors={"codes": codes}, + ... ) + >>> builder.fit(samples) + >>> builder.input_processors["codes"] is codes + True + >>> type(builder.input_processors["age"]).__name__ + 'TensorProcessor' """ def __init__( @@ -61,10 +84,10 @@ def __init__( self.input_schema = input_schema self.output_schema = output_schema self._input_processors = ( - input_processors if input_processors is not None else {} + dict(input_processors) if input_processors is not None else {} ) self._output_processors = ( - output_processors if output_processors is not None else {} + dict(output_processors) if output_processors is not None else {} ) self._patient_to_index: Dict[str, List[int]] = {} self._record_to_index: Dict[str, List[int]] = {} @@ -145,9 +168,9 @@ def fit( - Builds `patient_to_index` and `record_to_index` mappings by recording the sample indices associated with `patient_id` and `record_id`/`visit_id` fields. - - Instantiates and fits input/output processors from the provided - schemas (unless pre-fitted processors were supplied to the - constructor). + - Instantiates and fits a processor for every schema field that + has none. Pre-fitted processors supplied to the constructor are + used as they are; any remaining fields still get fitted ones. """ # Validate the samples input_keys = set(self.input_schema.keys()) @@ -171,17 +194,17 @@ def fit( if record_id is not None: self._record_to_index.setdefault(record_id, []).append(i) - # Fit processors if they were not provided - if not self._input_processors: - for key, spec in self.input_schema.items(): - processor = self._get_processor_instance(spec) - processor.fit(samples, key) - self._input_processors[key] = processor - if not self._output_processors: - for key, spec in self.output_schema.items(): - processor = self._get_processor_instance(spec) - processor.fit(samples, key) - self._output_processors[key] = processor + # Fit a processor for every schema field that has none. Supplied + # (pre-fitted) processors are used as they are, never refitted. + for schema, processors in ( + (self.input_schema, self._input_processors), + (self.output_schema, self._output_processors), + ): + for key, spec in schema.items(): + if key not in processors: + processor = self._get_processor_instance(spec) + processor.fit(samples, key) + processors[key] = processor self._fitted = True diff --git a/tests/core/test_partial_processors.py b/tests/core/test_partial_processors.py new file mode 100644 index 000000000..30b104915 --- /dev/null +++ b/tests/core/test_partial_processors.py @@ -0,0 +1,85 @@ +"""Supplying some processors must still fit the others. + +SampleBuilder used to fit input (or output) processors only when none were +supplied. Supplying a processor for one field left every other field of that +side without a processor, so its raw value (e.g. a Python list) reached the +model without any error. +""" + +import pickle +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset +from pyhealth.datasets.sample_dataset import SampleBuilder +from pyhealth.processors import ( + BinaryLabelProcessor, + SequenceProcessor, + TensorProcessor, +) + + +def _samples(): + return [ + {"patient_id": "p1", "codes": ["a", "b"], "age": [40.0], "label": 1}, + {"patient_id": "p2", "codes": ["c"], "age": [55.0], "label": 0}, + ] + + +class TestPartialProcessors(unittest.TestCase): + def setUp(self): + self.codes = SequenceProcessor() + self.codes.fit(_samples(), "codes") + + def test_missing_input_processor_is_fitted(self): + builder = SampleBuilder( + input_schema={"codes": "sequence", "age": "tensor"}, + output_schema={"label": "binary"}, + input_processors={"codes": self.codes}, + ) + builder.fit(_samples()) + self.assertIs(builder.input_processors["codes"], self.codes) + self.assertIsInstance(builder.input_processors["age"], TensorProcessor) + out = builder.transform({"sample": pickle.dumps(_samples()[0])}) + self.assertIsInstance(out["age"], torch.Tensor) + + def test_supplied_processor_is_not_refitted(self): + vocab = dict(self.codes.code_vocab) + samples = _samples() + [ + {"patient_id": "p3", "codes": ["z"], "age": [60.0], "label": 1} + ] + builder = SampleBuilder( + input_schema={"codes": "sequence", "age": "tensor"}, + output_schema={"label": "binary"}, + input_processors={"codes": self.codes}, + ) + builder.fit(samples) + self.assertEqual(self.codes.code_vocab, vocab) # "z" was not added + + def test_missing_output_processor_is_fitted(self): + samples = [dict(s, label2=s["label"]) for s in _samples()] + label = BinaryLabelProcessor() + label.fit(samples, "label") + builder = SampleBuilder( + input_schema={"codes": "sequence"}, + output_schema={"label": "binary", "label2": "binary"}, + output_processors={"label": label}, + ) + builder.fit(samples) + self.assertIs(builder.output_processors["label"], label) + self.assertIn("label2", builder.output_processors) + + def test_create_sample_dataset_with_partial_processors(self): + dataset = create_sample_dataset( + samples=_samples(), + input_schema={"codes": "sequence", "age": "tensor"}, + output_schema={"label": "binary"}, + input_processors={"codes": self.codes}, + dataset_name="test_partial_processors", + ) + self.assertIsInstance(dataset[0]["age"], torch.Tensor) + + +if __name__ == "__main__": + unittest.main() From 9ce56bad735f2416826d9787fbac1c626f1ba91b Mon Sep 17 00:00:00 2001 From: solarsys Date: Tue, 6 Oct 2026 02:07:37 -0400 Subject: [PATCH 2/6] Fit processors on the training split in set_task (split=PatientSplit) set_task fitted every processor (vocabularies, and any statistics a processor learns) on all samples, so validation and test patients shaped preprocessing. Splitting afterwards cannot refit. - New PatientSplit(ratios, seed): deterministic patient-level split (patients sorted by ID, then a seeded shuffle), 2 or 3 parts. - set_task(task, split=PatientSplit(...)) fits every processor on the training patients only and returns (train, val, test) or (train, test) as SampleDataset.subset views of one processed cache. - Streaming (S1): the split comes from patient_to_index, which fit() already builds; each processor then reads the same sequential sample stream through a masked view that skips other patients. No random access, nothing collected; the only new memory is the per-part int64 indices and a training mask (~9 bytes per sample). - Cache: split.npz holds the part indices; the processed-cache key adds the split spec only when given, so existing unsplit caches keep their key. schema.pkl records fit_split; SampleDataset.fit_split exposes it. - FeatureProcessor.learns_statistics (default False): when such a processor is fitted on all samples, SampleBuilder logs one warning pointing to split=. Vocabulary processors stay quiet. - tests/core/test_split_set_task.py (14): disjoint and complete parts, statistics from training patients (test shifted +100), test-only codes map to , fit_split recorded, cache reuse and per-seed caches, two-part split, unchanged unsplit cache key, supplied processors kept, the warning, invalid ratios, and fit streaming with a bounded peak. test_caching's signature check now includes split. - docs: splitter page section; processors page pointer. - examples/split_aware_set_task_mimic3demo.py and examples/benchmark_perf/benchmark_split_set_task.py (synthetic MIMIC-III at any scale, split vs no split, time and peak memory). Co-Authored-By: Claude Opus 5.5 --- .../datasets/pyhealth.datasets.splitter.rst | 37 +++ docs/api/processors.rst | 3 + .../benchmark_split_set_task.py | 235 ++++++++++++++++ examples/split_aware_set_task_mimic3demo.py | 62 +++++ pyhealth/datasets/__init__.py | 1 + pyhealth/datasets/base_dataset.py | 37 ++- pyhealth/datasets/sample_dataset.py | 70 ++++- pyhealth/datasets/splitter.py | 72 +++++ pyhealth/processors/base_processor.py | 23 ++ tests/core/test_caching.py | 2 + tests/core/test_split_set_task.py | 255 ++++++++++++++++++ 11 files changed, 792 insertions(+), 5 deletions(-) create mode 100644 examples/benchmark_perf/benchmark_split_set_task.py create mode 100644 examples/split_aware_set_task_mimic3demo.py create mode 100644 tests/core/test_split_set_task.py diff --git a/docs/api/datasets/pyhealth.datasets.splitter.rst b/docs/api/datasets/pyhealth.datasets.splitter.rst index d9e96447d..2b5924d61 100644 --- a/docs/api/datasets/pyhealth.datasets.splitter.rst +++ b/docs/api/datasets/pyhealth.datasets.splitter.rst @@ -3,6 +3,43 @@ Several data splitting function for `pyhealth.datasets` module to obtain training / validation / test sets. +Fitting processors on the training split +---------------------------------------- + +The ``split_by_*`` functions split an already-processed ``SampleDataset``: its +processors (code vocabularies, and any statistics a processor learns) were fitted +on all samples, validation and test patients included. To fit them on the +training patients only, let ``set_task`` do the split: + +.. code-block:: python + + from pyhealth.datasets import PatientSplit + + train, val, test = dataset.set_task( + task, split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42) + ) + train.fit_split # {'kind': 'patient', 'ratios': [0.7, 0.1, 0.2], 'seed': 42} + +- Every processor is fitted on the training patients; validation and test samples + are processed with those frozen processors, so a code seen only in test maps to + ````, as it would at deployment. +- Patients are sorted by ID before the seeded shuffle, so a ``PatientSplit`` gives + the same patients on every run and machine. Two ratios give ``(train, test)``. +- Samples are streamed while fitting: the training patients are selected from the + index of samples per patient that processing already keeps, and each processor + reads the same sample stream, skipping other patients. Memory does not grow with + the size of the samples. ``examples/benchmark_perf/benchmark_split_set_task.py`` + measures it at scale on synthetic data. +- The result is cached per split: a rerun with the same ``PatientSplit`` reuses it, + another seed or ratio builds its own. +- Processors that learn statistics from the data (``learns_statistics = True``) + log a warning when fitted on all samples, since that leaks validation/test + information. Vocabulary-only processors do not warn. + +The default, ``set_task(task)`` without ``split``, is unchanged: it fits on all +samples and returns one dataset. + + .. automodule:: pyhealth.datasets.splitter :members: :undoc-members: diff --git a/docs/api/processors.rst b/docs/api/processors.rst index 783e34353..ea83659b4 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -301,6 +301,9 @@ In PyHealth 2.0.2 and earlier, supplying a processor for some fields fitted none for the rest of that schema: those fields reached the model as raw Python values, without an error. See ``examples/reuse_train_processors.py``. +For datasets built with ``set_task``, ``set_task(task, split=PatientSplit(...))`` +does this in one step and streams the samples; see :doc:`datasets/pyhealth.datasets.splitter`. + Processor String Keys --------------------- diff --git a/examples/benchmark_perf/benchmark_split_set_task.py b/examples/benchmark_perf/benchmark_split_set_task.py new file mode 100644 index 000000000..d4496b758 --- /dev/null +++ b/examples/benchmark_perf/benchmark_split_set_task.py @@ -0,0 +1,235 @@ +"""Benchmark set_task with and without split=PatientSplit(...) at scale. + +Generates synthetic MIMIC-III-format tables (no real patient data) for N +patients, then runs MIMIC3Dataset + MortalityPredictionMIMIC3 twice in separate +processes, each with a fresh cache: + + nosplit : dataset.set_task(task) (current behaviour) + split : dataset.set_task(task, split=PatientSplit(...)) (fit on train only) + +For each run it reports wall time per stage and the peak resident memory of the +process plus its workers, sampled every 0.1 s. Use it to check that fitting on +the training split streams the samples (no extra memory) on your machine. + +Usage: + python examples/benchmark_perf/benchmark_split_set_task.py --patients 50000 + python examples/benchmark_perf/benchmark_split_set_task.py --patients 300000 --num-workers 4 + +Data and caches go under --workdir (default: a temporary folder, deleted at the +end; pass a path to keep them). Generation is reused when the folder exists. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +import tempfile +import threading +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import psutil + + +# ---------------------------------------------------------------- data ---- +def generate_mimic3(root: Path, n_patients: int, seed: int = 0) -> dict: + """Writes synthetic MIMIC-III CSVs with the columns PyHealth reads.""" + rng = np.random.default_rng(seed) + root.mkdir(parents=True, exist_ok=True) + subject = np.arange(10_000, 10_000 + n_patients) + + pd.DataFrame( + { + "row_id": np.arange(n_patients), + "subject_id": subject, + "gender": rng.choice(["F", "M"], n_patients), + "dob": "2080-01-01 00:00:00", + "dod": "", + "dod_hosp": "", + "dod_ssn": "", + "expire_flag": 0, + } + ).to_csv(root / "PATIENTS.csv", index=False) + + n_adm = rng.integers(1, 7, n_patients) # 1-6 admissions per patient + adm_subject = np.repeat(subject, n_adm) + n_a = adm_subject.size + hadm = np.arange(100_000, 100_000 + n_a) + order = np.concatenate([np.arange(k) for k in n_adm]) + admit = pd.Timestamp("2100-01-01") + pd.to_timedelta(order * 60 + rng.integers(0, 30, n_a), unit="D") + disch = admit + pd.to_timedelta(rng.integers(1, 15, n_a), unit="D") + pd.DataFrame( + { + "row_id": np.arange(n_a), "subject_id": adm_subject, "hadm_id": hadm, + "admittime": admit.strftime("%Y-%m-%d %H:%M:%S"), + "dischtime": disch.strftime("%Y-%m-%d %H:%M:%S"), + "deathtime": "", "admission_type": "EMERGENCY", + "admission_location": "EMERGENCY ROOM ADMIT", "discharge_location": "HOME", + "insurance": "Medicare", "language": "ENGL", "religion": "", "marital_status": "", + "ethnicity": "WHITE", "edregtime": "", "edouttime": "", "diagnosis": "", + "hospital_expire_flag": (rng.random(n_a) < 0.03).astype(int), + "has_chartevents_data": 1, + } + ).to_csv(root / "ADMISSIONS.csv", index=False) + + pd.DataFrame( + { + "subject_id": adm_subject, "intime": admit.strftime("%Y-%m-%d %H:%M:%S"), + "icustay_id": np.arange(200_000, 200_000 + n_a), "first_careunit": "MICU", + "dbsource": "carevue", "last_careunit": "MICU", + "outtime": disch.strftime("%Y-%m-%d %H:%M:%S"), + } + ).to_csv(root / "ICUSTAYS.csv", index=False) + + def codes_table(name, lo, hi, vocab, extra): + k = rng.integers(lo, hi + 1, n_a) + frame = pd.DataFrame( + { + "row_id": np.arange(k.sum()), + "subject_id": np.repeat(adm_subject, k), + "hadm_id": np.repeat(hadm, k), + **extra(k), + } + ) + frame.to_csv(root / name, index=False) + return int(k.sum()) + + icd = np.array([f"{c:05d}" for c in rng.choice(100_000, 6_000, replace=False)]) + proc = np.array([f"{c:04d}" for c in rng.choice(10_000, 2_000, replace=False)]) + ndc = np.array([f"{c:011d}" for c in rng.choice(10**11, 3_000, replace=False)]) + n_diag = codes_table( + "DIAGNOSES_ICD.csv", 3, 15, icd, + lambda k: {"seq_num": np.concatenate([np.arange(1, j + 1) for j in k]), + "icd9_code": rng.choice(icd, k.sum())}, + ) + n_proc = codes_table( + "PROCEDURES_ICD.csv", 1, 5, proc, + lambda k: {"seq_num": np.concatenate([np.arange(1, j + 1) for j in k]), + "icd9_code": rng.choice(proc, k.sum())}, + ) + + def presc_cols(k): + start = (np.repeat(admit, k)).strftime("%Y-%m-%d %H:%M:%S") + return { + "icustay_id": "", "startdate": start, "enddate": start, "drug_type": "MAIN", + "drug": "drug", "drug_name_poe": "", "drug_name_generic": "", "formulary_drug_cd": "", + "gsn": "", "ndc": rng.choice(ndc, k.sum()), "prod_strength": "", "dose_val_rx": "1", + "dose_unit_rx": "mg", "form_val_disp": "1", "form_unit_disp": "TAB", "route": "PO", + } + + n_presc = codes_table("PRESCRIPTIONS.csv", 2, 15, ndc, presc_cols) + return {"patients": n_patients, "admissions": int(n_a), "diagnoses": n_diag, + "procedures": n_proc, "prescriptions": n_presc} + + +# ------------------------------------------------------------- memory ---- +class PeakMemory: + """Peak RSS of this process plus its children, sampled in a thread.""" + + def __init__(self, interval: float = 0.1): + self.interval, self.peak, self._stop = interval, 0, threading.Event() + self._proc = psutil.Process() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _rss(self) -> int: + total = self._proc.memory_info().rss + for child in self._proc.children(recursive=True): + try: + total += child.memory_info().rss + except psutil.Error: + pass + return total + + def _run(self): + while not self._stop.is_set(): + self.peak = max(self.peak, self._rss()) + time.sleep(self.interval) + + def __enter__(self): + self.peak = self._rss() + self._thread.start() + return self + + def __exit__(self, *exc): + self._stop.set() + self._thread.join() + self.peak = max(self.peak, self._rss()) + + +# --------------------------------------------------------------- child ---- +def run_child(root: str, cache: str, mode: str, num_workers: int) -> dict: + from pyhealth.datasets import MIMIC3Dataset, PatientSplit + from pyhealth.tasks import MortalityPredictionMIMIC3 + + out = {"mode": mode, "baseline_rss_mib": psutil.Process().memory_info().rss / 2**20} + t0 = time.perf_counter() + with PeakMemory() as mem: + dataset = MIMIC3Dataset( + root=root, tables=["diagnoses_icd", "procedures_icd", "prescriptions"], + cache_dir=cache, num_workers=num_workers, + ) + _ = dataset.global_event_df + out["events_s"], out["events_peak_mib"] = time.perf_counter() - t0, mem.peak / 2**20 + + split = PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42) if mode == "split" else None + t1 = time.perf_counter() + with PeakMemory() as mem: + result = dataset.set_task(MortalityPredictionMIMIC3(), num_workers=num_workers, split=split) + out["set_task_s"], out["set_task_peak_mib"] = time.perf_counter() - t1, mem.peak / 2**20 + + parts = result if isinstance(result, tuple) else (result,) + out["samples"] = [len(p) for p in parts] + out["patients_in_index"] = len(parts[0].patient_to_index) if mode == "nosplit" else None + schema = Path(parts[0].path) / "schema.pkl" + out["schema_pkl_mib"] = schema.stat().st_size / 2**20 + split_file = Path(parts[0].path) / "split.npz" + out["split_npz_mib"] = split_file.stat().st_size / 2**20 if split_file.exists() else 0.0 + return out + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--patients", type=int, default=20_000) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--workdir", type=str, default=None) + parser.add_argument("--modes", type=str, default="nosplit,split") + parser.add_argument("--child", nargs=4, metavar=("ROOT", "CACHE", "MODE", "WORKERS")) + args = parser.parse_args() + + if args.child: + root, cache, mode, workers = args.child + print("RESULT " + json.dumps(run_child(root, cache, mode, int(workers))), flush=True) + return + + workdir = Path(args.workdir) if args.workdir else Path(tempfile.mkdtemp(prefix="pyhealth_split_bench_")) + data = workdir / f"mimic3_synthetic_{args.patients}" + try: + if not (data / "PRESCRIPTIONS.csv").exists(): + t = time.perf_counter() + sizes = generate_mimic3(data, args.patients) + print(f"generated {sizes} in {time.perf_counter() - t:.0f} s", flush=True) + for mode in args.modes.split(","): + cache = workdir / f"cache_{args.patients}_{mode}" + shutil.rmtree(cache, ignore_errors=True) + proc = subprocess.run( + [sys.executable, __file__, "--child", str(data), str(cache), mode, str(args.num_workers)], + capture_output=True, text=True, + ) + result = [l for l in proc.stdout.splitlines() if l.startswith("RESULT ")] + if proc.returncode or not result: + print(proc.stderr[-3000:], file=sys.stderr) + raise SystemExit(f"{mode} run failed (exit {proc.returncode})") + print(result[-1][7:], flush=True) + finally: + if not args.workdir: + shutil.rmtree(workdir, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/examples/split_aware_set_task_mimic3demo.py b/examples/split_aware_set_task_mimic3demo.py new file mode 100644 index 000000000..8f5f5f0e8 --- /dev/null +++ b/examples/split_aware_set_task_mimic3demo.py @@ -0,0 +1,62 @@ +"""Fit processors on training patients only with set_task(task, split=...). + +Splitting after set_task() means the code vocabularies (and any statistics a +processor learns) were fitted on every patient, test patients included. +Passing a PatientSplit to set_task fits them on the training patients only and +returns train / val / test datasets ready for a model. + +Uses the small MIMIC-III demo bundled in ``test-resources/``; no download. + +Usage (from the repository root): + python examples/split_aware_set_task_mimic3demo.py +""" + +import tempfile +from pathlib import Path + +from pyhealth.datasets import MIMIC3Dataset, PatientSplit, get_dataloader +from pyhealth.models import RNN +from pyhealth.tasks import MortalityPredictionMIMIC3 +from pyhealth.trainer import Trainer + +DEMO_ROOT = ( + Path(__file__).resolve().parent.parent / "test-resources" / "core" / "mimic3demo" +) + + +def main(): + with tempfile.TemporaryDirectory() as cache: + dataset = MIMIC3Dataset( + root=str(DEMO_ROOT), + tables=["diagnoses_icd", "procedures_icd", "prescriptions"], + cache_dir=cache, + ) + train, val, test = dataset.set_task( + MortalityPredictionMIMIC3(), + split=PatientSplit(ratios=(0.6, 0.2, 0.2), seed=0), + ) + # Ratios split patients, not samples: patients have different numbers + # of admissions, so sample counts need not follow the ratios. + print("patients (train, val, test):", + *(len(part.patient_to_index) for part in (train, val, test))) + print("samples (train, val, test):", len(train), len(val), len(test)) + print("processors fitted on:", train.fit_split) + + vocab = train.input_processors["conditions"].code_vocab + test_codes = {c for i in range(len(test)) for c in test[i]["conditions"].tolist()} + unknown = vocab[""] + print(f"training vocabulary: {len(vocab)} codes; " + f"test samples use : {unknown in test_codes}") + + model = RNN(dataset=train) + trainer = Trainer(model=model, metrics=["roc_auc", "pr_auc"], enable_logging=False) + trainer.train( + get_dataloader(train, batch_size=16, shuffle=True), + val_dataloader=get_dataloader(val, batch_size=16), + epochs=1, + ) + print("test:", trainer.evaluate(get_dataloader(test, batch_size=16))) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 57a9956c2..74062b8c8 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -71,6 +71,7 @@ def __init__(self, *args, **kwargs): from .support2 import Support2Dataset from .tcga_prad import TCGAPRADDataset from .splitter import ( + PatientSplit as PatientSplit, sample_balanced, split_by_patient, split_by_patient_conformal, diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 8534778a0..ad806e830 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -45,6 +45,7 @@ from ..processors.base_processor import FeatureProcessor from .configs import load_yaml_config from .sample_dataset import SampleDataset, SampleBuilder +from .splitter import PatientSplit from ..utils import set_env # Set logging level for distributed to ERROR to reduce verbosity @@ -1106,7 +1107,8 @@ def set_task( num_workers: Optional[int] = None, input_processors: Optional[Dict[str, FeatureProcessor]] = None, output_processors: Optional[Dict[str, FeatureProcessor]] = None, - ) -> SampleDataset: + split: PatientSplit | None = None, + ) -> SampleDataset | tuple[SampleDataset, ...]: """Processes the base dataset to generate the task-specific sample dataset. The cache structure is as follows:: @@ -1114,6 +1116,7 @@ def set_task( task_df.ld/ # Intermediate task dataframe based on schema samples_{proc_uuid}.ld/ # Final processed samples after applying processors schema.pkl # Saved SampleBuilder schema + split.npz # Sample indices per split part (with split= only) *.bin # Processed sample files Args: @@ -1125,9 +1128,23 @@ def set_task( output_processors (Optional[Dict[str, FeatureProcessor]]): Pre-fitted output processors. If provided, these will be used instead of creating new ones from task's output_schema. Defaults to None. + split (Optional[PatientSplit]): Split the samples by patient and fit + every processor on the training patients only, so validation and + test patients never shape preprocessing. Samples are streamed; + nothing is loaded into memory beyond the per-sample index that + processing already keeps. Defaults to None: fit on all samples + and return one dataset, as before. Returns: - SampleDataset: The generated sample dataset. + SampleDataset: The generated sample dataset, or, with ``split``, a + tuple of datasets (train, test) or (train, val, test) whose + processors were fitted on the training part. + + Examples: + >>> from pyhealth.datasets import PatientSplit + >>> train, val, test = dataset.set_task( # doctest: +SKIP + ... task, split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42) + ... ) Raises: AssertionError: If no default task is found and task is None. @@ -1180,6 +1197,8 @@ def set_task( if output_processors else None ), + # Only present with a split, so existing unsplit caches keep their key. + **({"split": split.to_dict()} if split is not None else {}), }, sort_keys=True, default=str, @@ -1240,8 +1259,12 @@ def _is_valid_litdata_cache(path: Path) -> bool: input_processors=input_processors, output_processors=output_processors, ) - builder.fit(dataset) + builder.fit(dataset, split=split) builder.save(str(samples_path / "schema.pkl")) + if split is not None: + # Written before the samples, whose index.json marks the + # cache complete, so a valid cache always has its split. + np.savez(samples_path / "split.npz", *builder.split_indices) # Apply processors and save final samples to cache_dir logger.info(f"Processing samples and saving to {samples_path}...") @@ -1256,11 +1279,17 @@ def _is_valid_litdata_cache(path: Path) -> bool: f"Found cached processed samples at {samples_path}, skipping processing." ) - return SampleDataset( + samples = SampleDataset( path=str(samples_path), dataset_name=self.dataset_name, task_name=task.task_name, ) + if split is None: + return samples + with np.load(samples_path / "split.npz") as parts: + return tuple( + samples.subset(parts[f"arr_{k}"]) for k in range(len(parts.files)) + ) def _main_guard(self, func_name: str): """Warn if method is accessed from a non-main process.""" diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 6a6813003..2816ceebd 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -10,10 +10,33 @@ import litdata from litdata.utilities.train_test_split import deepcopy_dataset import copy +import logging + +import numpy as np from ..processors import get_processor, IgnoreProcessor from ..processors.base_processor import FeatureProcessor +logger = logging.getLogger(__name__) + + +class _MaskedSamples: + """A re-iterable view of a sample stream that keeps only masked positions. + + Each iteration walks the underlying stream in order and yields the samples + whose position is True in ``mask``. Nothing is collected, so fitting on a + split costs no more memory than fitting on all samples. + """ + + def __init__(self, samples: Iterable[dict[str, Any]], mask: np.ndarray): + self._samples = samples + self._mask = mask + + def __iter__(self): + for keep, sample in zip(self._mask, self._samples): + if keep: + yield sample + def _remap_index_mapping( mapping: dict[str, list[int]], indices: Sequence[int] @@ -92,6 +115,8 @@ def __init__( self._patient_to_index: Dict[str, List[int]] = {} self._record_to_index: Dict[str, List[int]] = {} self._fitted = False + self._split = None + self._split_indices: list[np.ndarray] | None = None @property def input_processors(self) -> Dict[str, FeatureProcessor]: @@ -109,6 +134,15 @@ def output_processors(self) -> Dict[str, FeatureProcessor]: ) return self._output_processors + @property + def split_indices(self) -> list[np.ndarray] | None: + """Sample indices of each split part (train first), or None without a split.""" + if not self._fitted: + raise RuntimeError( + "SampleBuilder.fit must be called before accessing split_indices." + ) + return self._split_indices + @property def patient_to_index(self) -> Dict[str, List[int]]: if not self._fitted: @@ -153,6 +187,7 @@ def _get_processor_instance(self, processor_spec): def fit( self, samples: Iterable[Dict[str, Any]], + split: Any | None = None, ) -> None: """Fit processors and build mapping indices from an iterator of samples. @@ -161,6 +196,12 @@ def fit( sample should contain keys covering both the configured `input_schema` and `output_schema`. These samples are not required to be pickled; `fit` operates on in-memory dicts. + It is iterated several times and never collected, so a + streaming dataset works without loading it into memory. + split: Optional :class:`~pyhealth.datasets.PatientSplit`. When + given, processors are fitted on the first (training) part only, + and the sample indices of every part are kept in + ``split_indices``. Behavior: - Validates the samples contain all keys specified by the input @@ -186,7 +227,9 @@ def fit( # Build index mappings self._patient_to_index = {} self._record_to_index = {} + n_samples = 0 for i, sample in enumerate(samples): + n_samples = i + 1 patient_id = sample.get("patient_id") if patient_id is not None: self._patient_to_index.setdefault(patient_id, []).append(i) @@ -194,8 +237,20 @@ def fit( if record_id is not None: self._record_to_index.setdefault(record_id, []).append(i) + # With a split, fit on the training patients only: a masked view of + # the same stream, so nothing is loaded into memory. + self._split = split + self._split_indices = None + fit_samples = samples + if split is not None: + self._split_indices = split.split_indices(self._patient_to_index) + train_mask = np.zeros(n_samples, dtype=bool) + train_mask[self._split_indices[0]] = True + fit_samples = _MaskedSamples(samples, train_mask) + # Fit a processor for every schema field that has none. Supplied # (pre-fitted) processors are used as they are, never refitted. + learned_on_all = [] for schema, processors in ( (self.input_schema, self._input_processors), (self.output_schema, self._output_processors), @@ -203,8 +258,17 @@ def fit( for key, spec in schema.items(): if key not in processors: processor = self._get_processor_instance(spec) - processor.fit(samples, key) + processor.fit(fit_samples, key) processors[key] = processor + if split is None and getattr(processor, "learns_statistics", False): + learned_on_all.append(key) + if learned_on_all: + logger.warning( + "Processors for %s learn statistics from the data and were fitted " + "on all samples, so validation/test patients shaped them. Fit on " + "the training patients with set_task(task, split=PatientSplit(...)).", + ", ".join(learned_on_all), + ) self._fitted = True @@ -262,6 +326,7 @@ def save(self, path: str) -> None: "output_processors": self._output_processors, "patient_to_index": self._patient_to_index, "record_to_index": self._record_to_index, + "fit_split": self._split.to_dict() if self._split is not None else None, } with open(path, "wb") as f: pickle.dump(metadata, f) @@ -365,6 +430,9 @@ def __init__( self.patient_to_index = metadata["patient_to_index"] self.record_to_index = metadata["record_to_index"] + # The split the processors were fitted on, or None if fitted on all + # samples (also None for caches written before this was recorded). + self.fit_split: dict | None = metadata.get("fit_split") def _remove_ignored_processors(self): """Remove any processors that are IgnoreProcessor instances.""" diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 2dbc94186..df74fd2ec 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -1,3 +1,6 @@ +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from itertools import chain from typing import List, Optional, Tuple, Union @@ -10,6 +13,75 @@ # TODO: add more splitting methods +@dataclass(frozen=True) +class PatientSplit: + """A patient-level split for ``BaseDataset.set_task(task, split=...)``. + + With a split, ``set_task`` fits every processor on the training patients + only and returns one dataset per part, so validation and test patients + never shape preprocessing (vocabularies, learned statistics). + + Patients are sorted by ID before shuffling with ``seed``, so the same + ``PatientSplit`` gives the same patients on every run and machine, + whatever order the samples are stored in. + + Args: + ratios: Two or three fractions summing to 1: (train, test) or + (train, val, test). + seed: Seed for the shuffle. + + Examples: + >>> from pyhealth.datasets import PatientSplit + >>> split = PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42) + >>> parts = split.split_indices({"p1": [0, 1], "p2": [2], "p3": [3, 4]}) + >>> sorted(int(i) for part in parts for i in part) + [0, 1, 2, 3, 4] + """ + + ratios: Sequence[float] + seed: int = 0 + + def __post_init__(self): + ratios = tuple(float(r) for r in self.ratios) + if len(ratios) not in (2, 3): + raise ValueError("PatientSplit needs 2 or 3 ratios: (train, test) or (train, val, test).") + if any(r < 0 for r in ratios) or not math.isclose(sum(ratios), 1.0, abs_tol=1e-6): + raise ValueError(f"PatientSplit ratios must be non-negative and sum to 1, got {ratios}.") + object.__setattr__(self, "ratios", ratios) + + def to_dict(self) -> dict: + """JSON-serialisable description, used in cache keys and metadata.""" + return {"kind": "patient", "ratios": list(self.ratios), "seed": int(self.seed)} + + def split_indices( + self, patient_to_index: Mapping[str, Sequence[int]] + ) -> list[np.ndarray]: + """Splits sample indices by patient. + + Args: + patient_to_index: Patient ID to the indices of its samples. + + Returns: + One sorted ``int64`` array of sample indices per ratio, in order + (train first). Each patient's samples land in exactly one part. + """ + patients = sorted(patient_to_index, key=str) + order = np.random.default_rng(self.seed).permutation(len(patients)) + cuts = [ + int(len(patients) * sum(self.ratios[: k + 1])) + for k in range(len(self.ratios) - 1) + ] + parts = [] + for part in np.split(order, cuts): + indices = np.fromiter( + chain.from_iterable(patient_to_index[patients[p]] for p in part), + dtype=np.int64, + ) + indices.sort() + parts.append(indices) + return parts + + def _label_to_int(label) -> int: """Convert a stored label (int/np scalar/torch scalar) to Python int.""" if torch.is_tensor(label): diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 6f8f5ed63..650c143d1 100644 --- a/pyhealth/processors/base_processor.py +++ b/pyhealth/processors/base_processor.py @@ -63,8 +63,31 @@ class FeatureProcessor(Processor): Processor for individual fields (features). Example: Tokenization, image loading, normalization. + + Set ``learns_statistics = True`` on processors whose ``fit()`` learns values + from the data (means, standard deviations, medians, bins). Fitting such a + processor on all samples leaks validation/test information, so PyHealth + warns unless it is fitted on a training split + (``set_task(task, split=PatientSplit(...))``). + + Examples: + >>> import torch + >>> from pyhealth.processors.base_processor import FeatureProcessor + >>> class MeanCenter(FeatureProcessor): + ... learns_statistics = True + ... def fit(self, samples, field): + ... values = [s[field] for s in samples] + ... self.mean = sum(values) / len(values) + ... def process(self, value): + ... return torch.tensor(value - self.mean) + >>> p = MeanCenter() + >>> p.fit([{"x": 1.0}, {"x": 3.0}], "x") + >>> p.process(5.0) + tensor(3.) """ + learns_statistics: bool = False + def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: """Fit the processor to the samples. diff --git a/tests/core/test_caching.py b/tests/core/test_caching.py index a346b98e1..65556ab0d 100644 --- a/tests/core/test_caching.py +++ b/tests/core/test_caching.py @@ -123,6 +123,7 @@ def test_set_task_signature(self): "num_workers", "input_processors", "output_processors", + "split", ] self.assertEqual(params, expected_params) @@ -131,6 +132,7 @@ def test_set_task_signature(self): self.assertEqual(sig.parameters["num_workers"].default, None) self.assertEqual(sig.parameters["input_processors"].default, None) self.assertEqual(sig.parameters["output_processors"].default, None) + self.assertEqual(sig.parameters["split"].default, None) def test_set_task_writes_cache_and_metadata(self): """Ensure set_task materializes cache files and schema metadata.""" diff --git a/tests/core/test_split_set_task.py b/tests/core/test_split_set_task.py new file mode 100644 index 000000000..e1c6de565 --- /dev/null +++ b/tests/core/test_split_set_task.py @@ -0,0 +1,255 @@ +"""set_task(task, split=PatientSplit(...)): fit processors on training patients only. + +The synthetic data makes leakage visible: test patients' values are shifted by ++100 and they carry a code that never occurs in training. +""" + +import json +import logging +import tempfile +import tracemalloc +import unittest +import uuid +from pathlib import Path + +import numpy as np +import torch + +from pyhealth.datasets import BaseDataset, PatientSplit +from pyhealth.datasets.sample_dataset import SampleBuilder +from pyhealth.processors import SequenceProcessor +from pyhealth.processors.base_processor import FeatureProcessor +from pyhealth.tasks import BaseTask + +N_PATIENTS = 40 +SPLIT = PatientSplit(ratios=(0.6, 0.2, 0.2), seed=3) + +CONFIG = """version: "1.0" +tables: + events: + file_path: "events.csv" + patient_id: "patient_id" + timestamp: "time" + attributes: + - "code" + - "value" +""" + + +class MeanCenter(FeatureProcessor): + """Learns the mean of a numeric feature and subtracts it (a statistic).""" + + learns_statistics = True + + def __init__(self): + self.mean = None + + def fit(self, samples, field): + values = [float(s[field][0]) for s in samples] + self.mean = sum(values) / len(values) + + def process(self, value): + return torch.tensor([float(value[0]) - self.mean]) + + +class EventTask(BaseTask): + task_name = "split_test_task" + input_schema = {"codes": "sequence", "value": MeanCenter} + output_schema = {"label": "binary"} + + def __call__(self, patient): + samples = [] + for i, event in enumerate(patient.get_events("events")): + samples.append( + { + "patient_id": patient.patient_id, + "visit_id": f"{patient.patient_id}-{i}", + "codes": event.code.split("|"), + "value": [float(event.value)], + "label": i % 2 if len(patient.get_events("events")) > 1 else int(patient.patient_id[-1]) % 2, + } + ) + return samples + + +def _patient_ids(): + return [f"p{i:03d}" for i in range(N_PATIENTS)] + + +def _parts(split): + """Patients in each split part, from one-sample-per-patient indices.""" + ids = _patient_ids() + parts = split.split_indices({pid: [k] for k, pid in enumerate(ids)}) + return [{ids[int(k)] for k in part} for part in parts] + + +def _write_dataset(root: Path, test_patients: set) -> None: + rows = ["patient_id,time,code,value"] + for k, pid in enumerate(_patient_ids()): + for v in range(1 + k % 3): # 1-3 samples per patient + shift = 100.0 if pid in test_patients else 0.0 + code = "A|TESTONLY" if pid in test_patients else ("A|B" if v % 2 else "A") + rows.append(f"{pid},2020-01-{1 + v:02d} 00:00:00,{code},{10.0 + v + shift}") + (root / "events.csv").write_text("\n".join(rows) + "\n") + (root / "config.yaml").write_text(CONFIG) + + +class TestSplitSetTask(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.tmp = tempfile.TemporaryDirectory() + root = Path(cls.tmp.name) / "data" + root.mkdir() + cls.parts = _parts(SPLIT) + _write_dataset(root, test_patients=cls.parts[2]) + cls.dataset = BaseDataset( + root=str(root), + tables=["events"], + dataset_name="SplitDataset", + config_path=str(root / "config.yaml"), + cache_dir=str(Path(cls.tmp.name) / "cache"), + num_workers=1, + ) + cls.train, cls.val, cls.test = cls.dataset.set_task(EventTask(), split=SPLIT) + + @classmethod + def tearDownClass(cls): + cls.tmp.cleanup() + + @staticmethod + def _patients(ds): + return {ds[i]["patient_id"] for i in range(len(ds))} + + def test_parts_are_patient_disjoint_and_complete(self): + parts = [self._patients(d) for d in (self.train, self.val, self.test)] + self.assertEqual(parts, self.parts) + all_samples = self.dataset.set_task(EventTask()) + self.assertEqual(sum(len(d) for d in (self.train, self.val, self.test)), len(all_samples)) + + def test_statistics_come_from_training_patients(self): + train_values = [ + 10.0 + v + for k, pid in enumerate(_patient_ids()) + if pid in self.parts[0] + for v in range(1 + k % 3) + ] + mean = self.train.input_processors["value"].mean + self.assertAlmostEqual(mean, float(np.mean(train_values))) + self.assertLess(mean, 20.0) # fitted on all samples it would include +100 values + + def test_test_only_codes_map_to_unknown(self): + codes = self.train.input_processors["codes"] + self.assertNotIn("TESTONLY", codes.code_vocab) + sample = self.test[0] + self.assertIn(codes.code_vocab[""], sample["codes"].tolist()) + + def test_fit_split_is_recorded(self): + for part in (self.train, self.val, self.test): + self.assertEqual(part.fit_split, SPLIT.to_dict()) + self.assertIsNone(self.dataset.set_task(EventTask()).fit_split) + + def test_rerun_reuses_the_cache_and_the_same_split(self): + with self.assertLogs("pyhealth.datasets.base_dataset", "INFO") as logs: + again = self.dataset.set_task(EventTask(), split=SPLIT) + self.assertTrue(any("Found cached processed samples" in m for m in logs.output)) + self.assertFalse(any("Fitting processors" in m for m in logs.output)) + self.assertEqual([self._patients(d) for d in again], self.parts) + + def test_other_seed_gives_other_split_and_cache(self): + other = PatientSplit(ratios=(0.6, 0.2, 0.2), seed=4) + train, _, _ = self.dataset.set_task(EventTask(), split=other) + self.assertEqual(self._patients(train), _parts(other)[0]) + self.assertNotEqual(train.path, self.train.path) + + def test_two_part_split(self): + parts = self.dataset.set_task(EventTask(), split=PatientSplit(ratios=(0.8, 0.2))) + self.assertEqual(len(parts), 2) + + def test_unsplit_path_and_cache_key_are_unchanged(self): + samples = self.dataset.set_task(EventTask()) + proc_key = json.dumps( + {"input_processors": None, "output_processors": None}, sort_keys=True, default=str + ) + self.assertTrue(samples.path.endswith(f"samples_{uuid.uuid5(uuid.NAMESPACE_DNS, proc_key)}.ld")) + + def test_supplied_processor_is_kept_with_split(self): + codes = SequenceProcessor() + codes.fit([{"codes": ["A", "B", "Z"]}], "codes") + vocab = dict(codes.code_vocab) + train, _, test = self.dataset.set_task( + EventTask(), input_processors={"codes": codes}, split=SPLIT + ) + self.assertEqual(train.input_processors["codes"].code_vocab, vocab) + self.assertLess(train.input_processors["value"].mean, 20.0) + + +class TestSplitBuilder(unittest.TestCase): + def _samples(self, n=6): + return [ + {"patient_id": f"p{i}", "codes": ["A"], "value": [float(i)], "label": i % 2} + for i in range(n) + ] + + def test_warns_when_statistics_are_fitted_on_all_samples(self): + builder = SampleBuilder( + input_schema={"codes": "sequence", "value": MeanCenter}, + output_schema={"label": "binary"}, + ) + with self.assertLogs("pyhealth.datasets.sample_dataset", "WARNING") as logs: + builder.fit(self._samples()) + self.assertIn("value", "\n".join(logs.output)) + + def test_no_warning_with_a_split(self): + builder = SampleBuilder( + input_schema={"codes": "sequence", "value": MeanCenter}, + output_schema={"label": "binary"}, + ) + with self.assertNoLogs("pyhealth.datasets.sample_dataset", "WARNING"): + builder.fit(self._samples(), split=PatientSplit(ratios=(0.5, 0.5))) + + def test_no_warning_for_vocabulary_only(self): + builder = SampleBuilder( + input_schema={"codes": "sequence"}, output_schema={"label": "binary"} + ) + with self.assertNoLogs("pyhealth.datasets.sample_dataset", "WARNING"): + builder.fit(self._samples()) + + def test_split_fitting_streams_without_collecting_samples(self): + n, payload = 400, 4000 # ~50 MB if every sample were held at once + + class LazyStream: + def __iter__(self): + for i in range(n): + yield { + "patient_id": f"p{i % 100}", + "codes": ["A", "B"], + "value": [float(i)], + "label": i % 2, + "payload": [float(i)] * payload, + } + + def __len__(self): + raise AssertionError("fit must not take len() of the stream") + + builder = SampleBuilder( + input_schema={"codes": "sequence", "value": MeanCenter}, + output_schema={"label": "binary"}, + ) + tracemalloc.start() + builder.fit(LazyStream(), split=PatientSplit(ratios=(0.7, 0.3), seed=0)) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + self.assertLess(peak, 5 * 2**20, f"fit peaked at {peak / 2**20:.1f} MiB") + train_idx = builder.split_indices[0] + self.assertEqual(sum(len(p) for p in builder.split_indices), n) + expected = float(np.mean([float(i) for i in train_idx])) + self.assertAlmostEqual(builder.input_processors["value"].mean, expected) + + def test_invalid_ratios(self): + for ratios in ((0.5, 0.6), (1.0,), (0.5, -0.1, 0.6), (0.2, 0.2, 0.2, 0.4)): + with self.subTest(ratios=ratios), self.assertRaises(ValueError): + PatientSplit(ratios=ratios) + + +if __name__ == "__main__": + unittest.main() From 381c6a15cd251149238bc0143bc3fb703541644c Mon Sep 17 00:00:00 2001 From: John Wu Date: Tue, 6 Oct 2026 12:43:45 -0500 Subject: [PATCH 3/6] Merge LitData worker indexes in numeric rank order _proc_transform gives each worker a contiguous range of task samples, and patient_to_index / record_to_index are built in task order before processing, so they assume processed sample i is task sample i. _litdata_merge delegated to litdata's BinaryWriter.merge, which concatenates the per-worker "{rank}.index.json" files in string order (0, 1, 10, 11, ..., 19, 2, ...). With more than 10 workers that reorders the worker blocks, so patient_to_index points at other patients' samples and split_by_patient / set_task(split=...) return parts that share patients. On full MIMIC-IV mortality with num_workers=32, 120,222 of 128,236 indices were wrong and ~3.6k patients appeared in both train and val. Merge the indexes ourselves, ordered by integer rank. Adds a unit test for the merge with 12 ranks and an end-to-end 12-worker set_task test; both fail on the old merge. Caches built with more than 10 workers before this fix keep the wrong order and are still reused, because the cache key does not change. Co-Authored-By: Claude Opus 5.5 --- pyhealth/datasets/base_dataset.py | 36 ++++++++++++++++-- tests/core/test_litdata_merge.py | 62 +++++++++++++++++++++++++++++++ tests/core/test_split_set_task.py | 42 +++++++++++++++++++++ 3 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tests/core/test_litdata_merge.py diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index ad806e830..4e9c6df05 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -16,6 +16,7 @@ import multiprocessing import multiprocessing.queues import shutil +import time from filelock import FileLock import litdata @@ -206,6 +207,13 @@ def _litdata_merge(cache_dir: Path) -> None: """ Merges LitData binary writer index files in the given cache directory. + Workers write contiguous sample ranges, each with its own + ``{rank}.index.json``. The chunks are concatenated in numeric rank order so + that merged sample ``i`` is still input sample ``i``; ``patient_to_index`` + and ``record_to_index`` are built before processing and rely on this. + litdata's ``BinaryWriter.merge`` sorts the file names as strings + (0, 1, 10, 11, ..., 2), which reorders samples with more than 10 workers. + Args: cache_dir (Path): The cache directory containing LitData binary writer files. """ @@ -217,7 +225,10 @@ def _litdata_merge(cache_dir: Path) -> None: if _INDEX_FILENAME in files: return - index_files = [f for f in files if f.endswith(_INDEX_FILENAME)] + index_files = sorted( + (f for f in files if f.endswith(f".{_INDEX_FILENAME}")), + key=lambda f: int(f.split(".", 1)[0]), + ) # Return if there are no index files to merge if len(index_files) == 0: @@ -225,9 +236,26 @@ def _litdata_merge(cache_dir: Path) -> None: "There are zero samples in the dataset, please check the task and processors." ) - BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB").merge( - num_workers=len(index_files) - ) + chunks, config = [], None + for name in index_files: + with open(Path(cache_dir) / name) as f: + data = json.load(f) + if config is not None and data["config"] != config: + raise ValueError( + f"Inconsistent LitData config across worker index files in {cache_dir}." + ) + config = data["config"] + chunks.extend(data["chunks"]) + + # index.json marks the cache complete, so write it before removing the parts. + with open(Path(cache_dir) / _INDEX_FILENAME, "w") as f: + json.dump( + {"chunks": chunks, "config": config, "updated_at": str(time.time())}, + f, + sort_keys=True, + ) + for name in index_files: + os.remove(Path(cache_dir) / name) class _ProgressContext: diff --git a/tests/core/test_litdata_merge.py b/tests/core/test_litdata_merge.py new file mode 100644 index 000000000..5b190aad3 --- /dev/null +++ b/tests/core/test_litdata_merge.py @@ -0,0 +1,62 @@ +"""_litdata_merge must keep samples in worker order for any worker count. + +Workers write contiguous sample ranges, and patient_to_index (built before +processing) assumes merged sample i is input sample i. Sorting the per-worker +index files as strings puts rank 10 before rank 2, so this uses 12 ranks. +""" + +import tempfile +import unittest +from pathlib import Path + +import litdata +from litdata.streaming.writer import BinaryWriter + +from pyhealth.datasets.base_dataset import _litdata_merge +from pyhealth.utils import set_env + +N_WORKERS = 12 +PER_WORKER = 3 + + +def _write_worker_outputs(cache_dir: Path) -> None: + """Writes N_WORKERS contiguous ranges, as _proc_transform_fn does.""" + with set_env(DATA_OPTIMIZER_NUM_WORKERS=str(N_WORKERS)): + for rank in range(N_WORKERS): + with set_env(DATA_OPTIMIZER_GLOBAL_RANK=str(rank)): + writer = BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB") + for k in range(PER_WORKER): + i = rank * PER_WORKER + k + writer.add_item(k, {"position": i, "patient_id": f"p{i:03d}"}) + writer.done() + + +class TestLitdataMerge(unittest.TestCase): + def test_merge_keeps_worker_order_beyond_ten_workers(self): + with tempfile.TemporaryDirectory() as tmp: + cache = Path(tmp) + _write_worker_outputs(cache) + _litdata_merge(cache) + positions = [s["position"] for s in litdata.StreamingDataset(str(cache))] + self.assertEqual(positions, list(range(N_WORKERS * PER_WORKER))) + self.assertEqual( + sorted(p.name for p in cache.glob("*index.json")), ["index.json"] + ) + + def test_merge_is_noop_when_index_exists(self): + with tempfile.TemporaryDirectory() as tmp: + cache = Path(tmp) + _write_worker_outputs(cache) + _litdata_merge(cache) + before = (cache / "index.json").read_text() + _litdata_merge(cache) + self.assertEqual((cache / "index.json").read_text(), before) + + def test_merge_without_worker_indexes_raises(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(ValueError): + _litdata_merge(Path(tmp)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_split_set_task.py b/tests/core/test_split_set_task.py index e1c6de565..0a18dd029 100644 --- a/tests/core/test_split_set_task.py +++ b/tests/core/test_split_set_task.py @@ -251,5 +251,47 @@ def test_invalid_ratios(self): PatientSplit(ratios=ratios) +class TestSplitSetTaskManyWorkers(unittest.TestCase): + """More than 10 workers: merged samples must stay in task order. + + patient_to_index is built before processing, so a merge that reorders the + worker blocks (10.index.json before 2.index.json) hands every part the wrong samples. + """ + + N_WORKERS = 12 + + @classmethod + def setUpClass(cls): + cls.tmp = tempfile.TemporaryDirectory() + root = Path(cls.tmp.name) / "data" + root.mkdir() + cls.parts = _parts(SPLIT) + _write_dataset(root, test_patients=cls.parts[2]) + cls.dataset = BaseDataset( + root=str(root), + tables=["events"], + dataset_name="SplitDatasetManyWorkers", + config_path=str(root / "config.yaml"), + cache_dir=str(Path(cls.tmp.name) / "cache"), + num_workers=1, + ) + + @classmethod + def tearDownClass(cls): + cls.tmp.cleanup() + + def test_split_parts_hold_their_patients(self): + parts = self.dataset.set_task(EventTask(), num_workers=self.N_WORKERS, split=SPLIT) + got = [{d[i]["patient_id"] for i in range(len(d))} for d in parts] + self.assertEqual(got, self.parts) + + def test_patient_to_index_points_at_own_samples(self): + samples = self.dataset.set_task(EventTask(), num_workers=self.N_WORKERS) + self.assertGreater(len(samples), self.N_WORKERS * 2) + for pid, indices in samples.patient_to_index.items(): + for i in indices: + self.assertEqual(samples[i]["patient_id"], pid) + + if __name__ == "__main__": unittest.main() From 1c867623d9f98f8704157209a24142d548575174 Mon Sep 17 00:00:00 2001 From: John Wu Date: Tue, 6 Oct 2026 16:48:20 -0500 Subject: [PATCH 4/6] Use litdata's merge with litdata >=0.2.64,<0.2.68 instead of a custom merge The previous commit fixed the worker-index merge order by reimplementing BinaryWriter.merge in _litdata_merge. litdata 0.2.64 already sorts the per-worker "{rank}.index.json" files numerically, so go back to calling BinaryWriter.merge and require litdata >= 0.2.64 instead; this drops our copy of litdata's index format. Cap below 0.2.68: from 0.2.68 on, iterating SampleDataset.subset() parts (what split_by_patient and split= return) on a cache built with 0.2.68+ is ~10-35x slower (MIMIC-IV train part of a 0.2.76-built cache: 28-35 s with 0.2.68-0.2.76 vs 1.5-3.1 s with 0.2.64-0.2.67). The regression tests from the previous commit are unchanged: they fail with litdata 0.2.59 and pass with 0.2.64, 0.2.67 and 0.2.76. On full MIMIC-IV with num_workers=32 and litdata 0.2.67, unmodified master and this branch both give 0 misplaced samples and 0 patients shared between parts. pixi.lock is not updated here: it is already out of date with pyproject.toml, so re-locking also moves torch, dask and the CUDA wheels; that is a separate PR. Co-Authored-By: Claude Opus 5.5 --- pyhealth/datasets/base_dataset.py | 40 ++++++++----------------------- pyproject.toml | 4 +++- 2 files changed, 13 insertions(+), 31 deletions(-) diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 4e9c6df05..c5b592a2d 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -16,7 +16,6 @@ import multiprocessing import multiprocessing.queues import shutil -import time from filelock import FileLock import litdata @@ -208,11 +207,12 @@ def _litdata_merge(cache_dir: Path) -> None: Merges LitData binary writer index files in the given cache directory. Workers write contiguous sample ranges, each with its own - ``{rank}.index.json``. The chunks are concatenated in numeric rank order so - that merged sample ``i`` is still input sample ``i``; ``patient_to_index`` - and ``record_to_index`` are built before processing and rely on this. - litdata's ``BinaryWriter.merge`` sorts the file names as strings - (0, 1, 10, 11, ..., 2), which reorders samples with more than 10 workers. + ``{rank}.index.json``. ``patient_to_index`` and ``record_to_index`` are + built before processing and assume merged sample ``i`` is input sample + ``i``, so the per-worker indexes must be joined in numeric rank order. + ``BinaryWriter.merge`` does this from litdata 0.2.64 on; earlier versions + sort the file names as strings (0, 1, 10, 11, ..., 2), which reorders + samples with more than 10 workers. Args: cache_dir (Path): The cache directory containing LitData binary writer files. @@ -225,10 +225,7 @@ def _litdata_merge(cache_dir: Path) -> None: if _INDEX_FILENAME in files: return - index_files = sorted( - (f for f in files if f.endswith(f".{_INDEX_FILENAME}")), - key=lambda f: int(f.split(".", 1)[0]), - ) + index_files = [f for f in files if f.endswith(_INDEX_FILENAME)] # Return if there are no index files to merge if len(index_files) == 0: @@ -236,26 +233,9 @@ def _litdata_merge(cache_dir: Path) -> None: "There are zero samples in the dataset, please check the task and processors." ) - chunks, config = [], None - for name in index_files: - with open(Path(cache_dir) / name) as f: - data = json.load(f) - if config is not None and data["config"] != config: - raise ValueError( - f"Inconsistent LitData config across worker index files in {cache_dir}." - ) - config = data["config"] - chunks.extend(data["chunks"]) - - # index.json marks the cache complete, so write it before removing the parts. - with open(Path(cache_dir) / _INDEX_FILENAME, "w") as f: - json.dump( - {"chunks": chunks, "config": config, "updated_at": str(time.time())}, - f, - sort_keys=True, - ) - for name in index_files: - os.remove(Path(cache_dir) / name) + BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB").merge( + num_workers=len(index_files) + ) class _ProgressContext: diff --git a/pyproject.toml b/pyproject.toml index 49a2d73d5..62fe5750d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,9 @@ dependencies = [ "pandas~=2.3.1", "pydantic~=2.11.7", "dask[complete]>=2026.1.0", - "litdata~=0.2.59", + # >=0.2.64: worker indexes merge in numeric rank order (patient_to_index relies + # on it). <0.2.68: 0.2.68+ reads SampleDataset.subset() parts ~10-35x slower. + "litdata>=0.2.64,<0.2.68", "pyarrow>=23.0.1", "orjson~=3.10", "narwhals~=2.13.0", From 0c63fd4ab1536a91f1b84dedb654ec19b4c359a5 Mon Sep 17 00:00:00 2001 From: John Wu Date: Tue, 6 Oct 2026 17:38:02 -0500 Subject: [PATCH 5/6] Add a Split interface for set_task(split=...); allow any number of parts set_task(task, split=...) uses two methods of the split: split_indices() to get the sample indices of each part (part 0 is the training part the processors are fitted on) and to_dict() for the cache key and fit_split. Define them on an abstract Split class so custom splits have a documented contract; PatientSplit is now a Split and the reference implementation. set_task uses the parts exactly as returned: they may overlap or leave samples out. PatientSplit accepted only 2 or 3 ratios, although nothing downstream needs that. It now takes two or more ratios, training part first, e.g. (train, val, cal, test). The ratios must still be non-negative and sum to 1; each check has its own error message. Tests: a 4-part PatientSplit, a custom held-out-patients Split through set_task (test-only code maps to ), and a custom Split with overlapping parts that set_task returns as given. Docs gain a "Writing your own split" section. Co-Authored-By: Claude Opus 5.5 --- .../datasets/pyhealth.datasets.splitter.rst | 39 ++++++++++- pyhealth/datasets/__init__.py | 1 + pyhealth/datasets/base_dataset.py | 14 ++-- pyhealth/datasets/sample_dataset.py | 13 ++-- pyhealth/datasets/splitter.py | 69 +++++++++++++++++-- tests/core/test_split_set_task.py | 64 ++++++++++++++++- 6 files changed, 180 insertions(+), 20 deletions(-) diff --git a/docs/api/datasets/pyhealth.datasets.splitter.rst b/docs/api/datasets/pyhealth.datasets.splitter.rst index 2b5924d61..899b8ba14 100644 --- a/docs/api/datasets/pyhealth.datasets.splitter.rst +++ b/docs/api/datasets/pyhealth.datasets.splitter.rst @@ -24,7 +24,9 @@ training patients only, let ``set_task`` do the split: are processed with those frozen processors, so a code seen only in test maps to ````, as it would at deployment. - Patients are sorted by ID before the seeded shuffle, so a ``PatientSplit`` gives - the same patients on every run and machine. Two ratios give ``(train, test)``. + the same patients on every run and machine. It takes two or more ratios summing + to 1, training part first: ``(train, test)``, ``(train, val, test)``, + ``(train, val, cal, test)`` and so on, returned in that order. - Samples are streamed while fitting: the training patients are selected from the index of samples per patient that processing already keeps, and each processor reads the same sample stream, skipping other patients. Memory does not grow with @@ -39,6 +41,41 @@ training patients only, let ``set_task`` do the split: The default, ``set_task(task)`` without ``split``, is unchanged: it fits on all samples and returns one dataset. +Writing your own split +---------------------- + +``split=`` accepts any subclass of ``Split``. It implements two methods: + +- ``split_indices(patient_to_index)`` returns one array of sample indices per part, + training part first. ``patient_to_index`` maps each patient ID to the indices of + its samples. ``set_task`` fits the processors on part 0 and returns every part as + given: parts may overlap or leave samples out. +- ``to_dict()`` returns a JSON-serialisable description. It goes into the cache key, + so it must change whenever the parts would, and is saved as ``fit_split``. + +For example, holding out a fixed set of patients: + +.. code-block:: python + + import numpy as np + from pyhealth.datasets import Split + + class HoldoutSplit(Split): + def __init__(self, test_patients): + self.test_patients = sorted(test_patients) + + def split_indices(self, patient_to_index): + test = set(self.test_patients) + parts = ([], []) + for pid, indices in patient_to_index.items(): + parts[pid in test].extend(indices) + return [np.array(sorted(p), dtype=np.int64) for p in parts] + + def to_dict(self): + return {"kind": "holdout", "test_patients": self.test_patients} + + train, test = dataset.set_task(task, split=HoldoutSplit(["patient-a", "patient-b"])) + .. automodule:: pyhealth.datasets.splitter :members: diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 74062b8c8..7acd921e6 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -72,6 +72,7 @@ def __init__(self, *args, **kwargs): from .tcga_prad import TCGAPRADDataset from .splitter import ( PatientSplit as PatientSplit, + Split as Split, sample_balanced, split_by_patient, split_by_patient_conformal, diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index c5b592a2d..9a16c668f 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -45,7 +45,7 @@ from ..processors.base_processor import FeatureProcessor from .configs import load_yaml_config from .sample_dataset import SampleDataset, SampleBuilder -from .splitter import PatientSplit +from .splitter import Split from ..utils import set_env # Set logging level for distributed to ERROR to reduce verbosity @@ -1115,7 +1115,7 @@ def set_task( num_workers: Optional[int] = None, input_processors: Optional[Dict[str, FeatureProcessor]] = None, output_processors: Optional[Dict[str, FeatureProcessor]] = None, - split: PatientSplit | None = None, + split: Split | None = None, ) -> SampleDataset | tuple[SampleDataset, ...]: """Processes the base dataset to generate the task-specific sample dataset. The cache structure is as follows:: @@ -1136,16 +1136,18 @@ def set_task( output_processors (Optional[Dict[str, FeatureProcessor]]): Pre-fitted output processors. If provided, these will be used instead of creating new ones from task's output_schema. Defaults to None. - split (Optional[PatientSplit]): Split the samples by patient and fit - every processor on the training patients only, so validation and - test patients never shape preprocessing. Samples are streamed; + split (Optional[Split]): Split the samples, e.g. by patient with + ``PatientSplit``, and fit every processor on the first + (training) part only, so the other parts never shape + preprocessing. The parts are used as the split returns them. + Samples are streamed; nothing is loaded into memory beyond the per-sample index that processing already keeps. Defaults to None: fit on all samples and return one dataset, as before. Returns: SampleDataset: The generated sample dataset, or, with ``split``, a - tuple of datasets (train, test) or (train, val, test) whose + tuple with one dataset per part, training part first, whose processors were fitted on the training part. Examples: diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 2816ceebd..dfd99b2a2 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -4,6 +4,7 @@ import shutil import tempfile from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type +from typing import TYPE_CHECKING import inspect import random from bisect import bisect_right @@ -17,6 +18,9 @@ from ..processors import get_processor, IgnoreProcessor from ..processors.base_processor import FeatureProcessor +if TYPE_CHECKING: # splitter imports this module + from .splitter import Split + logger = logging.getLogger(__name__) @@ -187,7 +191,7 @@ def _get_processor_instance(self, processor_spec): def fit( self, samples: Iterable[Dict[str, Any]], - split: Any | None = None, + split: "Split | None" = None, ) -> None: """Fit processors and build mapping indices from an iterator of samples. @@ -198,9 +202,10 @@ def fit( required to be pickled; `fit` operates on in-memory dicts. It is iterated several times and never collected, so a streaming dataset works without loading it into memory. - split: Optional :class:`~pyhealth.datasets.PatientSplit`. When - given, processors are fitted on the first (training) part only, - and the sample indices of every part are kept in + split: Optional :class:`~pyhealth.datasets.Split`, such as + :class:`~pyhealth.datasets.PatientSplit`. When given, + processors are fitted on the first (training) part only, and + the sample indices of every part are kept in ``split_indices``. Behavior: diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index df74fd2ec..10eeb0c80 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -1,4 +1,5 @@ import math +from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass from itertools import chain @@ -13,8 +14,57 @@ # TODO: add more splitting methods +class Split(ABC): + """Interface for ``BaseDataset.set_task(task, split=...)``. + + ``set_task`` fits every processor on the samples of part 0 and returns one + ``SampleDataset`` per part, in order, as views over one processed cache. + It uses the parts exactly as returned: they may overlap or leave samples + out. :class:`PatientSplit` is the reference implementation. + + Examples: + >>> import numpy as np + >>> from pyhealth.datasets import Split + >>> class HoldoutSplit(Split): + ... def __init__(self, test_patients): + ... self.test_patients = sorted(test_patients) + ... def split_indices(self, patient_to_index): + ... test = set(self.test_patients) + ... parts = ([], []) + ... for pid, indices in patient_to_index.items(): + ... parts[pid in test].extend(indices) + ... return [np.array(sorted(p), dtype=np.int64) for p in parts] + ... def to_dict(self): + ... return {"kind": "holdout", "test_patients": self.test_patients} + >>> HoldoutSplit(["p2"]).split_indices({"p1": [0, 1], "p2": [2]}) + [array([0, 1]), array([2])] + """ + + @abstractmethod + def split_indices( + self, patient_to_index: Mapping[str, Sequence[int]] + ) -> list[np.ndarray]: + """Returns the sample indices of each part. + + Args: + patient_to_index: Patient ID to the indices of its samples. + + Returns: + One ``int64`` array of sample indices per part, training part + first. Sorted indices keep reads from the cache sequential. + """ + + @abstractmethod + def to_dict(self) -> dict: + """Returns a JSON-serialisable description of the split. + + It goes into the processed-cache key, so it must change whenever the + parts would, and is saved as ``SampleDataset.fit_split``. + """ + + @dataclass(frozen=True) -class PatientSplit: +class PatientSplit(Split): """A patient-level split for ``BaseDataset.set_task(task, split=...)``. With a split, ``set_task`` fits every processor on the training patients @@ -26,8 +76,9 @@ class PatientSplit: whatever order the samples are stored in. Args: - ratios: Two or three fractions summing to 1: (train, test) or - (train, val, test). + ratios: Two or more non-negative fractions summing to 1. The first + part is the training part, e.g. (train, test), (train, val, test) + or (train, val, cal, test). seed: Seed for the shuffle. Examples: @@ -43,10 +94,12 @@ class PatientSplit: def __post_init__(self): ratios = tuple(float(r) for r in self.ratios) - if len(ratios) not in (2, 3): - raise ValueError("PatientSplit needs 2 or 3 ratios: (train, test) or (train, val, test).") - if any(r < 0 for r in ratios) or not math.isclose(sum(ratios), 1.0, abs_tol=1e-6): - raise ValueError(f"PatientSplit ratios must be non-negative and sum to 1, got {ratios}.") + if len(ratios) < 2: + raise ValueError(f"PatientSplit needs at least 2 ratios, training part first; got {ratios}.") + if any(r < 0 for r in ratios): + raise ValueError(f"PatientSplit ratios must be non-negative, got {ratios}.") + if not math.isclose(sum(ratios), 1.0, abs_tol=1e-6): + raise ValueError(f"PatientSplit ratios must sum to 1, got {ratios} (sum {sum(ratios)}).") object.__setattr__(self, "ratios", ratios) def to_dict(self) -> dict: @@ -58,6 +111,8 @@ def split_indices( ) -> list[np.ndarray]: """Splits sample indices by patient. + Patients are sorted, shuffled with ``seed`` and cut by ``ratios``. + Args: patient_to_index: Patient ID to the indices of its samples. diff --git a/tests/core/test_split_set_task.py b/tests/core/test_split_set_task.py index 0a18dd029..7c968a9e3 100644 --- a/tests/core/test_split_set_task.py +++ b/tests/core/test_split_set_task.py @@ -15,7 +15,7 @@ import numpy as np import torch -from pyhealth.datasets import BaseDataset, PatientSplit +from pyhealth.datasets import BaseDataset, PatientSplit, Split from pyhealth.datasets.sample_dataset import SampleBuilder from pyhealth.processors import SequenceProcessor from pyhealth.processors.base_processor import FeatureProcessor @@ -72,6 +72,34 @@ def __call__(self, patient): return samples +class HoldoutSplit(Split): + """Custom split: the given patients form part 1, everyone else part 0.""" + + def __init__(self, test_patients): + self.test_patients = sorted(test_patients) + + def split_indices(self, patient_to_index): + test = set(self.test_patients) + parts = ([], []) + for pid, indices in patient_to_index.items(): + parts[pid in test].extend(indices) + return [np.array(sorted(p), dtype=np.int64) for p in parts] + + def to_dict(self): + return {"kind": "holdout", "test_patients": self.test_patients} + + +class OverlapSplit(Split): + """Custom split whose part 1 repeats part 0: set_task must accept it as given.""" + + def split_indices(self, patient_to_index): + everything = np.array(sorted(i for ix in patient_to_index.values() for i in ix), dtype=np.int64) + return [everything, everything[:5]] + + def to_dict(self): + return {"kind": "overlap"} + + def _patient_ids(): return [f"p{i:03d}" for i in range(N_PATIENTS)] @@ -165,6 +193,38 @@ def test_two_part_split(self): parts = self.dataset.set_task(EventTask(), split=PatientSplit(ratios=(0.8, 0.2))) self.assertEqual(len(parts), 2) + def test_four_part_split(self): + split = PatientSplit(ratios=(0.5, 0.1, 0.1, 0.3), seed=3) + parts = self.dataset.set_task(EventTask(), split=split) + self.assertEqual(len(parts), 4) + patients = [self._patients(d) for d in parts] + self.assertEqual(patients, _parts(split)) + for i in range(4): + for j in range(i + 1, 4): + self.assertFalse(patients[i] & patients[j]) + self.assertEqual(set().union(*patients), set(_patient_ids())) + self.assertEqual(sum(len(d) for d in parts), len(self.dataset.set_task(EventTask()))) + self.assertEqual(parts[3].fit_split["ratios"], [0.5, 0.1, 0.1, 0.3]) + + def test_custom_split(self): + split = HoldoutSplit(self.parts[2]) + train, held = self.dataset.set_task(EventTask(), split=split) + self.assertEqual(self._patients(held), self.parts[2]) + self.assertEqual(self._patients(train), set(_patient_ids()) - self.parts[2]) + self.assertEqual(held.fit_split, split.to_dict()) + codes = train.input_processors["codes"] + self.assertNotIn("TESTONLY", codes.code_vocab) + self.assertIn(codes.code_vocab[""], held[0]["codes"].tolist()) + + def test_custom_split_parts_are_used_as_given(self): + everything, first = self.dataset.set_task(EventTask(), split=OverlapSplit()) + self.assertEqual(len(everything), len(self.dataset.set_task(EventTask()))) + self.assertEqual(len(first), 5) + self.assertEqual( + [first[i]["visit_id"] for i in range(5)], + [everything[i]["visit_id"] for i in range(5)], + ) + def test_unsplit_path_and_cache_key_are_unchanged(self): samples = self.dataset.set_task(EventTask()) proc_key = json.dumps( @@ -246,7 +306,7 @@ def __len__(self): self.assertAlmostEqual(builder.input_processors["value"].mean, expected) def test_invalid_ratios(self): - for ratios in ((0.5, 0.6), (1.0,), (0.5, -0.1, 0.6), (0.2, 0.2, 0.2, 0.4)): + for ratios in ((0.5, 0.6), (1.0,), (0.5, -0.1, 0.6), (0.2, 0.2, 0.2, 0.3)): with self.subTest(ratios=ratios), self.assertRaises(ValueError): PatientSplit(ratios=ratios) From f3e65376d51791c0482499aec2f25e0da75942cc Mon Sep 17 00:00:00 2001 From: John Wu Date: Tue, 6 Oct 2026 19:12:06 -0500 Subject: [PATCH 6/6] Fix split_by_* docstrings; point them to set_task(split=...) The split_by_* functions return dataset.subset(...) views, which are SampleDatasets over the same processed samples, but their docstrings still described torch.utils.data.Subset objects with a .dataset attribute. Describe the actual return values, document get_index on split_by_sample, add a usage example to each (required by tools/check_pr_rules.py for modified public functions), and add a note that the dataset's processors were fitted on all samples, with a pointer to set_task(split=...) with PatientSplit or a custom Split. PatientSplit's docstring now says how it differs from split_by_patient: parts are chosen before fitting, and split_by_patient shuffles patients in dataset order, so it generally picks different patients for the same ratios and seed. Docstrings only; no behavior change. Co-Authored-By: Claude Opus 5.5 --- pyhealth/datasets/splitter.py | 180 +++++++++++++++++++++++++++------- 1 file changed, 147 insertions(+), 33 deletions(-) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 10eeb0c80..3f4678618 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -75,6 +75,12 @@ class PatientSplit(Split): ``PatientSplit`` gives the same patients on every run and machine, whatever order the samples are stored in. + Unlike :func:`split_by_patient`, which splits an already-processed + ``SampleDataset`` whose processors were fitted on all samples, the parts + are chosen before fitting. ``split_by_patient`` shuffles patients in the + order they appear in the dataset, so for the same ratios and seed it + generally picks different patients. + Args: ratios: Two or more non-negative fractions summing to 1. The first part is the training part, e.g. (train, test), (train, val, test) @@ -244,12 +250,21 @@ def split_by_visit( seed: random seed for shuffling the dataset Returns: - train_dataset, val_dataset, test_dataset: three subsets of the dataset of - type `torch.utils.data.Subset`. + train_dataset, val_dataset, test_dataset: ``SampleDataset`` views + (``dataset.subset``) over the same processed samples. + + Examples: + >>> from pyhealth.datasets import split_by_visit + >>> train, val, test = split_by_visit( # doctest: +SKIP + ... samples, [0.8, 0.1, 0.1], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, and `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert sum(ratios) == 1.0, "ratios must sum to 1.0" @@ -279,12 +294,21 @@ def split_by_patient( seed: random seed for shuffling the dataset Returns: - train_dataset, val_dataset, test_dataset: three subsets of the dataset of - type `torch.utils.data.Subset`. + train_dataset, val_dataset, test_dataset: ``SampleDataset`` views + (``dataset.subset``) over the same processed samples. + + Examples: + >>> from pyhealth.datasets import split_by_patient + >>> train, val, test = split_by_patient( # doctest: +SKIP + ... samples, [0.8, 0.1, 0.1], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, and `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert sum(ratios) == 1.0, "ratios must sum to 1.0" @@ -319,14 +343,26 @@ def split_by_sample( dataset: a `SampleDataset` object ratios: a list/tuple of ratios for train / val / test seed: random seed for shuffling the dataset + get_index: if True, return three ``torch.Tensor`` index vectors + instead of datasets Returns: - train_dataset, val_dataset, test_dataset: three subsets of the dataset of - type `torch.utils.data.Subset`. + train_dataset, val_dataset, test_dataset: ``SampleDataset`` views + (``dataset.subset``) over the same processed samples, or index + tensors if ``get_index`` is True. + + Examples: + >>> from pyhealth.datasets import split_by_sample + >>> train, val, test = split_by_sample( # doctest: +SKIP + ... samples, [0.8, 0.1, 0.1], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, and `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert sum(ratios) == 1.0, "ratios must sum to 1.0" @@ -364,13 +400,22 @@ def split_by_visit_conformal( seed: random seed for shuffling the dataset Returns: - train_dataset, val_dataset, cal_dataset, test_dataset: four subsets - of the dataset of type `torch.utils.data.Subset`. + train_dataset, val_dataset, cal_dataset, test_dataset: + ``SampleDataset`` views (``dataset.subset``) over the same + processed samples. + + Examples: + >>> from pyhealth.datasets import split_by_visit_conformal + >>> train, val, cal, test = split_by_visit_conformal( # doctest: +SKIP + ... samples, [0.6, 0.1, 0.1, 0.2], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, `cal_dataset.dataset`, and - `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test" @@ -410,13 +455,22 @@ def split_by_patient_conformal( seed: random seed for shuffling the dataset Returns: - train_dataset, val_dataset, cal_dataset, test_dataset: four subsets - of the dataset of type `torch.utils.data.Subset`. + train_dataset, val_dataset, cal_dataset, test_dataset: + ``SampleDataset`` views (``dataset.subset``) over the same + processed samples. + + Examples: + >>> from pyhealth.datasets import split_by_patient_conformal + >>> train, val, cal, test = split_by_patient_conformal( # doctest: +SKIP + ... samples, [0.6, 0.1, 0.1, 0.2], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, `cal_dataset.dataset`, and - `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test" @@ -481,10 +535,23 @@ def split_by_patient_conformal_tuh( to 1.0. seed: random seed used to shuffle the patient list. get_index: if ``True``, return four :class:`torch.Tensor` index vectors - instead of :class:`~torch.utils.data.Subset` objects. + instead of datasets. Returns: ``(train_dataset, val_dataset, cal_dataset, test_dataset)`` + + Examples: + >>> from pyhealth.datasets import split_by_patient_conformal_tuh + >>> train, val, cal, test = split_by_patient_conformal_tuh( # doctest: +SKIP + ... samples, [0.6, 0.2, 0.2], seed=42 + ... ) + + Note: + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ assert len(ratios) == 3, ( "ratios must have exactly 3 elements (train/val/cal). " @@ -556,10 +623,23 @@ def split_by_sample_conformal_tuh( ratios: the fraction of the train pool assigned to train / val / cal respectively seed: random seed for shuffling the train pool get_index: if True, return four ``torch.Tensor`` index vectors instead - of ``Subset`` objects + of datasets Returns: train_dataset, val_dataset, cal_dataset, test_dataset + + Examples: + >>> from pyhealth.datasets import split_by_sample_conformal_tuh + >>> train, val, cal, test = split_by_sample_conformal_tuh( # doctest: +SKIP + ... samples, [0.6, 0.2, 0.2], seed=42 + ... ) + + Note: + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ assert len(ratios) == 3, ( "ratios must have exactly 3 elements (train/val/cal). " @@ -638,10 +718,23 @@ def split_by_patient_tuh( to 1.0. seed: random seed used to shuffle the patient list. get_index: if ``True``, return three :class:`torch.Tensor` index - vectors instead of :class:`~torch.utils.data.Subset` objects. + vectors instead of datasets. Returns: ``(train_dataset, val_dataset, test_dataset)`` + + Examples: + >>> from pyhealth.datasets import split_by_patient_tuh + >>> train, val, test = split_by_patient_tuh( # doctest: +SKIP + ... samples, [0.8, 0.2], seed=42 + ... ) + + Note: + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ assert len(ratios) == 2, ( "ratios must have exactly 2 elements (train/val). " @@ -714,10 +807,23 @@ def split_by_sample_tuh( respectively. Must be a length-2 sequence summing to 1.0. seed: random seed for shuffling the train pool get_index: if True, return three ``torch.Tensor`` index vectors instead - of ``Subset`` objects + of datasets Returns: train_dataset, val_dataset, test_dataset + + Examples: + >>> from pyhealth.datasets import split_by_sample_tuh + >>> train, val, test = split_by_sample_tuh( # doctest: +SKIP + ... samples, [0.8, 0.2], seed=42 + ... ) + + Note: + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ assert len(ratios) == 2, ( "ratios must have exactly 2 elements (train/val). " @@ -776,17 +882,25 @@ def split_by_sample_conformal( dataset: a `SampleDataset` object ratios: a list/tuple of ratios for train / val / cal / test seed: random seed for shuffling the dataset - get_index: if True, return indices instead of Subset objects + get_index: if True, return indices instead of datasets Returns: - train_dataset, val_dataset, cal_dataset, test_dataset: four subsets - of the dataset of type `torch.utils.data.Subset`, or four tensors - of indices if get_index=True. + train_dataset, val_dataset, cal_dataset, test_dataset: + ``SampleDataset`` views (``dataset.subset``) over the same + processed samples, or four index tensors if ``get_index`` is True. + + Examples: + >>> from pyhealth.datasets import split_by_sample_conformal + >>> train, val, cal, test = split_by_sample_conformal( # doctest: +SKIP + ... samples, [0.6, 0.1, 0.1, 0.2], seed=42 + ... ) Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, `cal_dataset.dataset`, and - `test_dataset.dataset`. + The processors of ``dataset`` (vocabularies, learned statistics) were + fitted on all of its samples, so every part shaped them. To fit them on + the training part only, split inside ``set_task`` instead: + ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` + or a custom :class:`Split`. """ rng = np.random.default_rng(seed) assert len(ratios) == 4, "ratios must have 4 elements for train/val/cal/test"