From f9a60f6483b9ac6da5471919876c1bafc57c767f Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:21:38 +0530 Subject: [PATCH 01/13] Implement predefined dataset splitting function Added a new function to split datasets based on predefined train/test partitions, allowing for patient-based or sample-based splits with specified ratios. --- pyhealth/datasets/splitter.py | 287 +++++++++++++--------------------- 1 file changed, 105 insertions(+), 182 deletions(-) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index d3ee2109e..5ec2e1d78 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -509,6 +509,99 @@ def split_by_patient_conformal( return train_dataset, val_dataset, cal_dataset, test_dataset +def split_by_predefined( + dataset: SampleDataset, + ratios: Union[Tuple[float, ...], List[float]], + seed: Optional[int] = None, + get_index: Optional[bool] = False, + by_patient: bool = False, + split_key: str = "split", + train_value: str = "train", +): + """Splits a dataset that ships with an official train/test partition. + + Every sample must carry ``split_key``. Samples whose value equals + ``train_value`` form the pool that is shuffled and cut by ``ratios``; + all other samples are the held-out test set. This is the one place that + holds the logic for pre-defined splits: a dataset whose task tags each + sample with its official partition should call this function (directly or + through a thin wrapper) instead of adding a new splitter here. + + Args: + dataset: a ``SampleDataset``; needs ``patient_to_index`` if + ``by_patient=True``. + ratios: 2 values (train / val) or 3 values (train / val / cal), + summing to 1.0. Applied to the train pool only. + seed: random seed for shuffling the pool. With ``by_patient=True`` + patients are sorted by ID before shuffling, so the split does not + depend on the order of the samples. + get_index: if True, return ``torch.Tensor`` index vectors instead of + datasets. + by_patient: keep all of a patient's samples in the same split. The + first sample of each patient decides which partition the patient + belongs to. + split_key: sample field holding the official partition label. + train_value: value of ``split_key`` that marks the train pool. + + Returns: + ``(train, val, test)`` for 2 ratios, ``(train, val, cal, test)`` for 3 + ratios; ``SampleDataset`` views, or index tensors if ``get_index``. + + Examples: + >>> from pyhealth.datasets import split_by_predefined + >>> train, val, test = split_by_predefined( # doctest: +SKIP + ... samples, [0.8, 0.2], seed=42, by_patient=True + ... ) + """ + assert len(ratios) in (2, 3), ( + "ratios must have 2 (train/val) or 3 (train/val/cal) elements" + ) + assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" + + pool: dict = {} # group key (patient ID or sample index) -> sample indices + holdout: List[int] = [] + + if by_patient: + # A patient belongs to one partition, so the first sample decides. + for pid, indices in dataset.patient_to_index.items(): + first = dataset[indices[0]] + assert split_key in first, ( + f"Patient {pid}: sample missing '{split_key}' field." + ) + if first[split_key] == train_value: + pool[pid] = list(indices) + else: + holdout.extend(indices) + # Sort first: dict order follows sample order, which callers may change. + keys = sorted(pool, key=str) + else: + for i in range(len(dataset)): + sample = dataset[i] # read once per sample + assert split_key in sample, ( + f"Sample {i} is missing the '{split_key}' field." + ) + if sample[split_key] == train_value: + pool[i] = [i] + else: + holdout.append(i) + keys = list(pool) # already ascending sample indices + + rng = np.random.default_rng(seed) + rng.shuffle(keys) + + n = len(keys) + bounds = [0] + [int(n * c) for c in np.cumsum(ratios)[:-1]] + [n] + parts = [ + np.array(list(chain(*[pool[k] for k in keys[a:b]])), dtype=int) + for a, b in zip(bounds[:-1], bounds[1:]) + ] + parts.append(np.array(holdout, dtype=int)) + + if get_index: + return tuple(torch.tensor(p) for p in parts) + return tuple(dataset.subset(p) for p in parts) # type: ignore + + def split_by_patient_conformal_tuh( dataset: SampleDataset, ratios: Union[Tuple[float, float, float], List[float]], @@ -529,6 +622,8 @@ def split_by_patient_conformal_tuh( patients), never both. The test set is therefore always the TUH eval partition (fixed, fully held-out patients). + This is a thin wrapper around :func:`split_by_predefined`. + Args: dataset: a ``SampleDataset`` produced by ``EEGEventsTUEV`` or ``EEGAbnormalTUAB``. Each sample must carry a ``"split"`` field @@ -561,57 +656,7 @@ def split_by_patient_conformal_tuh( "ratios must have exactly 3 elements (train/val/cal). " "The test set is determined by the TUH eval partition." ) - assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" - - # Bucket patients by partition using dataset.patient_to_index (fast path). - # TUH guarantees each patient is in exactly one partition, so inspecting the - # first sample per patient is sufficient. - train_patient_to_indices: dict = {} - test_list: List[int] = [] - - for pid, indices in dataset.patient_to_index.items(): - first_sample = dataset[indices[0]] - assert "split" in first_sample, ( - f"Patient {pid}: sample missing 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - if first_sample["split"] == "train": - train_patient_to_indices[pid] = list(indices) - else: - test_list.extend(list(indices)) - - # Shuffle patients deterministically - patient_ids = sorted(train_patient_to_indices, key=str) - rng = np.random.default_rng(seed) - rng.shuffle(patient_ids) - - n = len(patient_ids) - train_end = int(n * ratios[0]) - val_end = int(n * (ratios[0] + ratios[1])) - - train_pids = patient_ids[:train_end] - val_pids = patient_ids[train_end:val_end] - cal_pids = patient_ids[val_end:] - - train_index = np.array(list(chain(*[train_patient_to_indices[p] for p in train_pids]))) - val_index = np.array(list(chain(*[train_patient_to_indices[p] for p in val_pids]))) - cal_index = np.array(list(chain(*[train_patient_to_indices[p] for p in cal_pids]))) - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(cal_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(cal_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) def split_by_sample_conformal_tuh( @@ -622,6 +667,8 @@ def split_by_sample_conformal_tuh( ): """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined train/eval split. + This is a thin wrapper around :func:`split_by_predefined`. + Args: dataset: a ``SampleDataset`` object produced by ``EEGEventsTUEV`` or ``EEGAbnormalTUAB`` ratios: the fraction of the train pool assigned to train / val / cal respectively @@ -649,52 +696,7 @@ def split_by_sample_conformal_tuh( "ratios must have exactly 3 elements (train/val/cal). " "The test set is determined by the dataset's own eval partition." ) - assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" - - # verify every sample has the required "split" field - for i in range(len(dataset)): - assert "split" in dataset[i], ( - f"Sample {i} is missing the 'split' field. " - "Make sure you used EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - - train_pool: List[int] = [] - test_list: List[int] = [] - for i in range(len(dataset)): - if dataset[i]["split"] == "train": - train_pool.append(i) - else: - test_list.append(i) - - # shuffle only the train pool - rng = np.random.default_rng(seed) - train_arr = np.array(train_pool) - rng.shuffle(train_arr) - - # Slice into train / val / cal. - n = len(train_arr) - train_end = int(n * ratios[0]) - val_end = int(n * (ratios[0] + ratios[1])) - - train_index = train_arr[:train_end] - val_index = train_arr[train_end:val_end] - cal_index = train_arr[val_end:] - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(cal_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(cal_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) def split_by_patient_tuh( @@ -712,6 +714,8 @@ def split_by_patient_tuh( Ensures that **all windows from a given patient stay in the same split** (train or val), preventing within-patient data leakage. + This is a thin wrapper around :func:`split_by_predefined`. + Args: dataset: a ``SampleDataset`` produced by ``EEGEventsTUEV`` or ``EEGAbnormalTUAB``. Each sample must carry a ``"split"`` field @@ -744,52 +748,7 @@ def split_by_patient_tuh( "ratios must have exactly 2 elements (train/val). " "The test set is determined by the TUH eval partition." ) - assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" - - train_patient_to_indices: dict = {} - test_list: List[int] = [] - - for pid, indices in dataset.patient_to_index.items(): - first_sample = dataset[indices[0]] - assert "split" in first_sample, ( - f"Patient {pid}: sample missing 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - if first_sample["split"] == "train": - train_patient_to_indices[pid] = list(indices) - else: - test_list.extend(list(indices)) - - patient_ids = sorted(train_patient_to_indices, key=str) - rng = np.random.default_rng(seed) - rng.shuffle(patient_ids) - - n = len(patient_ids) - train_end = int(n * ratios[0]) - - train_pids = patient_ids[:train_end] - val_pids = patient_ids[train_end:] - - train_index = np.array( - list(chain(*[train_patient_to_indices[p] for p in train_pids])) - ) - val_index = np.array( - list(chain(*[train_patient_to_indices[p] for p in val_pids])) - ) - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) def split_by_sample_tuh( @@ -804,6 +763,8 @@ def split_by_sample_tuh( (train / val / test) instead of 4-way (train / val / cal / test). The test set is always the TUH eval partition; no calibration set is produced. + This is a thin wrapper around :func:`split_by_predefined`. + Args: dataset: a ``SampleDataset`` object produced by ``EEGEventsTUEV`` or ``EEGAbnormalTUAB`` @@ -833,45 +794,7 @@ def split_by_sample_tuh( "ratios must have exactly 2 elements (train/val). " "The test set is determined by the dataset's own eval partition." ) - assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" - - for i in range(len(dataset)): - assert "split" in dataset[i], ( - f"Sample {i} is missing the 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - - train_pool: List[int] = [] - test_list: List[int] = [] - for i in range(len(dataset)): - if dataset[i]["split"] == "train": - train_pool.append(i) - else: - test_list.append(i) - - rng = np.random.default_rng(seed) - train_arr = np.array(train_pool) - rng.shuffle(train_arr) - - n = len(train_arr) - train_end = int(n * ratios[0]) - - train_index = train_arr[:train_end] - val_index = train_arr[train_end:] - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) def split_by_sample_conformal( From d42079944815f7ab0b13dc89d0899f9ceb23fcb5 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:25:59 +0530 Subject: [PATCH 02/13] Add split_by_predefined to dataset imports --- pyhealth/datasets/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 7acd921e6..87bccad57 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -78,6 +78,7 @@ def __init__(self, *args, **kwargs): split_by_patient_conformal, split_by_patient_conformal_tuh, split_by_patient_tuh, + split_by_predefined, split_by_sample, split_by_sample_conformal, split_by_sample_conformal_tuh, From 20482803811538627a27c6961131f0bd082ae05e Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:29:00 +0530 Subject: [PATCH 03/13] Include 'split' in returned sample structure Added 'split' field to the sample return structure. --- pyhealth/tasks/temple_university_EEG_tasks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyhealth/tasks/temple_university_EEG_tasks.py b/pyhealth/tasks/temple_university_EEG_tasks.py index 13e8c7206..122f99311 100644 --- a/pyhealth/tasks/temple_university_EEG_tasks.py +++ b/pyhealth/tasks/temple_university_EEG_tasks.py @@ -20,6 +20,7 @@ class EEGEventsTUEV(BaseTask): - "signal": np.ndarray, shape (16, 256*5) - "offending_channel": int - "label": int + - "split": str, "train" or "eval", the dataset's official partition. Used by pyhealth.datasets.split_by_predefined. Examples: >>> from pyhealth.datasets import TUEVDataset From fc6aeb7dee2cea68d4a9b975f6b0e4fecc4f6f17 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:30:49 +0530 Subject: [PATCH 04/13] Add tests for predefined dataset splitter --- tests/core/test_splitter_predefined.py | 49 ++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 tests/core/test_splitter_predefined.py diff --git a/tests/core/test_splitter_predefined.py b/tests/core/test_splitter_predefined.py new file mode 100644 index 000000000..6f9beafb4 --- /dev/null +++ b/tests/core/test_splitter_predefined.py @@ -0,0 +1,49 @@ +import numpy as np +import torch +from pyhealth.datasets.splitter import split_by_predefined + + +class FakeDataset: + def __init__(self, samples): + self.samples = samples + self.patient_to_index = {} + for i, s in enumerate(samples): + self.patient_to_index.setdefault(s["patient_id"], []).append(i) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, i): + return self.samples[i] + + def subset(self, indices): + return [int(i) for i in indices] + + +def make_ds(): + s = [{"patient_id": f"t{p}", "split": "train"} for p in range(10) for _ in range(3)] + s += [{"patient_id": f"e{p}", "split": "eval"} for p in range(4) for _ in range(2)] + return FakeDataset(s) + + +def test_test_set_is_eval_partition(): + ds = make_ds() + *_, test = split_by_predefined(ds, [0.8, 0.2], seed=0) + assert all(ds[i]["split"] == "eval" for i in test) and len(test) == 8 + + +def test_patient_level_has_no_leakage(): + ds = make_ds() + train, val, cal, test = split_by_predefined(ds, [0.6, 0.2, 0.2], seed=0, by_patient=True) + groups = [{ds[i]["patient_id"] for i in part} for part in (train, val, cal, test)] + for a in range(4): + for b in range(a + 1, 4): + assert groups[a].isdisjoint(groups[b]) + + +def test_seed_is_deterministic_and_get_index(): + ds = make_ds() + a = split_by_predefined(ds, [0.8, 0.2], seed=1, get_index=True) + b = split_by_predefined(ds, [0.8, 0.2], seed=1, get_index=True) + assert all(torch.equal(x, y) for x, y in zip(a, b)) + assert sum(len(p) for p in a) == len(ds) From 431d8f83926287517665126aee0e11fa183bacc2 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:32:21 +0530 Subject: [PATCH 05/13] Document dataset split porting process Added guidelines for porting official dataset splits to CONTRIBUTING.md. --- CONTRIBUTING.md | 49 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c2a2899ca..fcd55c6e3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -77,6 +77,53 @@ which runs `tools/check_pr_rules.py`. You can run the same check locally: python tools/check_pr_rules.py --base --head HEAD ``` +## Porting a Dataset's Official Splits + +Many datasets ship with pre-defined splits, for example the train/eval partitions +of the TUH EEG corpora (TUEV, TUAB). To keep `pyhealth/datasets/splitter.py` from +growing one near-identical function per dataset, port official splits as follows. + +1. **Tag every sample in the task.** In the task's `__call__`, set a `"split"` + field on each sample to the dataset's own partition label. `EEGEventsTUEV` and + `EEGAbnormalTUAB` in `pyhealth/tasks/temple_university_EEG_tasks.py` are the + reference: they iterate over the `"train"` and `"eval"` partitions and write + the partition name into `sample["split"]`. Document the field in the task's + docstring. +2. **Do not add a new splitter to `splitter.py`.** Call + `pyhealth.datasets.split_by_predefined`, which holds the logic for all + pre-defined splits: + - `ratios` has 2 values (train / val) or 3 values (train / val / cal) and is + applied to the train pool only. Samples whose `split_key` value is not + `train_value` form the held-out test set. + - `by_patient=True` keeps all of a patient's samples in the same split + (the dataset must have `patient_to_index`). Patients are sorted by ID before + shuffling, so the result does not depend on sample order. + - `split_key` (default `"split"`) and `train_value` (default `"train"`) adapt + the function to a different field name or partition label. + - `get_index=True` returns `torch.Tensor` index vectors instead of datasets. +3. **Wrap it only if the arguments are fixed.** If a dataset always needs the same + arguments, add a thin wrapper next to that dataset (in its module under + `pyhealth/datasets/`), not in `splitter.py`. The existing `*_tuh` functions in + `splitter.py` are such wrappers, kept for backward compatibility. +4. **Test it.** Add a unit test that checks the held-out set is exactly the + official eval partition, that no patient appears in more than one part when + `by_patient=True`, and that the same seed gives the same split. + +```python +>>> from pyhealth.datasets import split_by_predefined +>>> train, val, test = split_by_predefined( # doctest: +SKIP +... sample_dataset, [0.8, 0.2], seed=42, by_patient=True +... ) +>>> train, val, cal, test = split_by_predefined( # doctest: +SKIP +... sample_dataset, [0.6, 0.2, 0.2], seed=42, by_patient=True +... ) +``` + +These functions split an already-processed `SampleDataset`, so its processors +(vocabularies, learned statistics) were fitted on all samples. If preprocessing +must be fitted on the training part only, split inside `set_task` instead, with +`PatientSplit` or a custom `Split`. + ## Community We welcome you to join our community @@ -134,4 +181,4 @@ All code changes happen through pull requests. We enforce code reviews for the main branches. - `master`: requires two approvals from the core developers -- `develop`: requires one approval from the core developers \ No newline at end of file +- `develop`: requires one approval from the core developers From 71dc1f14c9226b5dbdc295ac042c1ee577a76568 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:30:45 +0530 Subject: [PATCH 06/13] Update splitter.py --- pyhealth/datasets/splitter.py | 39 +++++++++++++++++------------------ 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 15f4700f9..85bb11b38 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -727,12 +727,12 @@ def split_by_patient_conformal_tuh( ``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). " - "The test set is determined by the TUH eval partition." - ) - return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) + _validate_ratios(ratios, expected_length=3) + + + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) + # Bucket patients by partition using dataset.patient_to_index (fast path). # TUH guarantees each patient is in exactly one partition, so inspecting the @@ -791,6 +791,9 @@ def split_by_sample_conformal_tuh( seed: Optional[int] = None, get_index: Optional[bool] = False, ): + + _validate_ratios(ratios, expected_length=3) + """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined train/eval split. This is a thin wrapper around :func:`split_by_predefined`. @@ -827,12 +830,9 @@ def split_by_sample_conformal_tuh( ``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). " - "The test set is determined by the dataset's own eval partition." - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) - _validate_ratios(ratios, expected_length=3) + # verify every sample has the required "split" field for i in range(len(dataset)): @@ -886,6 +886,9 @@ def split_by_patient_tuh( seed: Optional[int] = None, get_index: Optional[bool] = False, ): + + _validate_ratios(ratios, expected_length=2) + """Splits a TUH EEG dataset by PATIENT within the TUH train partition. Like :func:`split_by_patient_conformal_tuh` but returns a 3-way split @@ -933,12 +936,8 @@ def split_by_patient_tuh( ``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). " - "The test set is determined by the TUH eval partition." - ) return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) - _validate_ratios(ratios, expected_length=2) + train_patient_to_indices: dict = {} test_list: List[int] = [] @@ -992,6 +991,9 @@ def split_by_sample_tuh( seed: Optional[int] = None, get_index: Optional[bool] = False, ): + + _validate_ratios(ratios, expected_length=2) + """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined split. Like :func:`split_by_sample_conformal_tuh` but returns a 3-way split @@ -1033,12 +1035,9 @@ def split_by_sample_tuh( ``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). " - "The test set is determined by the dataset's own eval partition." - ) + return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) - _validate_ratios(ratios, expected_length=2) + for i in range(len(dataset)): assert "split" in dataset[i], ( From ce70e20e7b6351b52619ec8ed528a30950446c3f Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:36:07 +0530 Subject: [PATCH 07/13] Update __init__.py --- pyhealth/datasets/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 0cadae3f1..af55a227a 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -79,7 +79,6 @@ def __init__(self, *args, **kwargs): split_by_patient_conformal, split_by_patient_conformal_tuh, split_by_patient_tuh, - split_by_predefined, split_by_sample, split_by_sample_conformal, split_by_sample_conformal_tuh, From d6dd21aa32ced1f2dace5a74c944b92c3dccd203 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:50:11 +0530 Subject: [PATCH 08/13] Update type hints in splitter.py Refactor type hints from Optional and Union to use the new PEP 604 syntax for type annotations. --- pyhealth/datasets/splitter.py | 266 +++++----------------------------- 1 file changed, 36 insertions(+), 230 deletions(-) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 85bb11b38..74f1240c3 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -4,7 +4,6 @@ from dataclasses import dataclass from itertools import chain from numbers import Real -from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -185,7 +184,7 @@ def sample_balanced( dataset: SampleDataset, ratio: float = 1.0, subsample: float = 1.0, - seed: Optional[int] = None, + seed: int | None = None, ) -> SampleDataset: """Keep positives and negatives at a target ratio, then cap total size. @@ -209,8 +208,8 @@ def sample_balanced( rng = np.random.default_rng(seed) - pos_indices: List[int] = [] - neg_indices: List[int] = [] + pos_indices: list[int] = [] + neg_indices: list[int] = [] for idx in range(len(dataset)): label = _label_to_int(dataset[idx]["label"]) @@ -270,8 +269,8 @@ def sample_balanced( def split_by_visit( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, + ratios: tuple[float, float, float] | list[float], + seed: int | None = None, ): """Splits the dataset by visit (i.e., samples). @@ -321,8 +320,8 @@ def split_by_visit( def split_by_patient( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, + ratios: tuple[float, float, float] | list[float], + seed: int | None = None, ): """Splits the dataset by patient. @@ -381,9 +380,9 @@ def split_by_patient( def split_by_sample( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): """Splits the dataset by sample @@ -444,8 +443,8 @@ def split_by_sample( def split_by_visit_conformal( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float, float], List[float]], - seed: Optional[int] = None, + ratios: tuple[float, float, float, float] | list[float], + seed: int | None = None, ): """Splits the dataset by visit (i.e., samples) for conformal prediction. @@ -506,8 +505,8 @@ def split_by_visit_conformal( def split_by_patient_conformal( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float, float], List[float]], - seed: Optional[int] = None, + ratios: tuple[float, float, float, float] | list[float], + seed: int | None = None, ): """Splits the dataset by patient for conformal prediction. @@ -577,9 +576,9 @@ def split_by_patient_conformal( def split_by_predefined( dataset: SampleDataset, - ratios: Union[Tuple[float, ...], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, ...] | list[float], + seed: int | None = None, + get_index: bool = False, by_patient: bool = False, split_key: str = "split", train_value: str = "train", @@ -625,7 +624,7 @@ def split_by_predefined( assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" pool: dict = {} # group key (patient ID or sample index) -> sample indices - holdout: List[int] = [] + holdout: list[int] = [] if by_patient: # A patient belongs to one partition, so the first sample decides. @@ -670,9 +669,9 @@ def split_by_predefined( def split_by_patient_conformal_tuh( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): """Splits a TUH EEG dataset by PATIENT within the TUH train partition. @@ -727,73 +726,16 @@ def split_by_patient_conformal_tuh( ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` or a custom :class:`Split`. """ - _validate_ratios(ratios, expected_length=3) - - return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) - - - # Bucket patients by partition using dataset.patient_to_index (fast path). - # TUH guarantees each patient is in exactly one partition, so inspecting the - # first sample per patient is sufficient. - train_patient_to_indices: dict = {} - test_list: List[int] = [] - - for pid, indices in dataset.patient_to_index.items(): - first_sample = dataset[indices[0]] - assert "split" in first_sample, ( - f"Patient {pid}: sample missing 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - if first_sample["split"] == "train": - train_patient_to_indices[pid] = list(indices) - else: - test_list.extend(list(indices)) - - # Shuffle patients deterministically - patient_ids = sorted(train_patient_to_indices, key=str) - rng = np.random.default_rng(seed) - rng.shuffle(patient_ids) - - n = len(patient_ids) - train_end = int(n * ratios[0]) - val_end = int(n * (ratios[0] + ratios[1])) - - train_pids = patient_ids[:train_end] - val_pids = patient_ids[train_end:val_end] - cal_pids = patient_ids[val_end:] - - train_index = np.array(list(chain(*[train_patient_to_indices[p] for p in train_pids]))) - val_index = np.array(list(chain(*[train_patient_to_indices[p] for p in val_pids]))) - cal_index = np.array(list(chain(*[train_patient_to_indices[p] for p in cal_pids]))) - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(cal_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(cal_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) def split_by_sample_conformal_tuh( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): - - _validate_ratios(ratios, expected_length=3) - """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined train/eval split. This is a thin wrapper around :func:`split_by_predefined`. @@ -830,65 +772,16 @@ def split_by_sample_conformal_tuh( ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` or a custom :class:`Split`. """ - + _validate_ratios(ratios, expected_length=3) return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) - - - # verify every sample has the required "split" field - for i in range(len(dataset)): - assert "split" in dataset[i], ( - f"Sample {i} is missing the 'split' field. " - "Make sure you used EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - - train_pool: List[int] = [] - test_list: List[int] = [] - for i in range(len(dataset)): - if dataset[i]["split"] == "train": - train_pool.append(i) - else: - test_list.append(i) - - # shuffle only the train pool - rng = np.random.default_rng(seed) - train_arr = np.array(train_pool) - rng.shuffle(train_arr) - - # Slice into train / val / cal. - n = len(train_arr) - train_end = int(n * ratios[0]) - val_end = int(n * (ratios[0] + ratios[1])) - - train_index = train_arr[:train_end] - val_index = train_arr[train_end:val_end] - cal_index = train_arr[val_end:] - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(cal_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(cal_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) def split_by_patient_tuh( dataset: SampleDataset, - ratios: Union[Tuple[float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): - - _validate_ratios(ratios, expected_length=2) - """Splits a TUH EEG dataset by PATIENT within the TUH train partition. Like :func:`split_by_patient_conformal_tuh` but returns a 3-way split @@ -936,64 +829,16 @@ def split_by_patient_tuh( ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` or a custom :class:`Split`. """ + _validate_ratios(ratios, expected_length=2) return split_by_predefined(dataset, ratios, seed, get_index, by_patient=True) - train_patient_to_indices: dict = {} - test_list: List[int] = [] - - for pid, indices in dataset.patient_to_index.items(): - first_sample = dataset[indices[0]] - assert "split" in first_sample, ( - f"Patient {pid}: sample missing 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - if first_sample["split"] == "train": - train_patient_to_indices[pid] = list(indices) - else: - test_list.extend(list(indices)) - - patient_ids = sorted(train_patient_to_indices, key=str) - rng = np.random.default_rng(seed) - rng.shuffle(patient_ids) - - n = len(patient_ids) - train_end = int(n * ratios[0]) - - train_pids = patient_ids[:train_end] - val_pids = patient_ids[train_end:] - - train_index = np.array( - list(chain(*[train_patient_to_indices[p] for p in train_pids])) - ) - val_index = np.array( - list(chain(*[train_patient_to_indices[p] for p in val_pids])) - ) - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) - - def split_by_sample_tuh( dataset: SampleDataset, - ratios: Union[Tuple[float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): - - _validate_ratios(ratios, expected_length=2) - """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined split. Like :func:`split_by_sample_conformal_tuh` but returns a 3-way split @@ -1035,54 +880,15 @@ def split_by_sample_tuh( ``base_dataset.set_task(task, split=...)`` with :class:`PatientSplit` or a custom :class:`Split`. """ - + _validate_ratios(ratios, expected_length=2) return split_by_predefined(dataset, ratios, seed, get_index, by_patient=False) - for i in range(len(dataset)): - assert "split" in dataset[i], ( - f"Sample {i} is missing the 'split' field. " - "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." - ) - - train_pool: List[int] = [] - test_list: List[int] = [] - for i in range(len(dataset)): - if dataset[i]["split"] == "train": - train_pool.append(i) - else: - test_list.append(i) - - rng = np.random.default_rng(seed) - train_arr = np.array(train_pool) - rng.shuffle(train_arr) - - n = len(train_arr) - train_end = int(n * ratios[0]) - - train_index = train_arr[:train_end] - val_index = train_arr[train_end:] - test_index = np.array(test_list) - - if get_index: - return ( - torch.tensor(train_index), - torch.tensor(val_index), - torch.tensor(test_index), - ) - else: - return ( - dataset.subset(train_index), # type: ignore - dataset.subset(val_index), # type: ignore - dataset.subset(test_index), # type: ignore - ) - - def split_by_sample_conformal( dataset: SampleDataset, - ratios: Union[Tuple[float, float, float, float], List[float]], - seed: Optional[int] = None, - get_index: Optional[bool] = False, + ratios: tuple[float, float, float, float] | list[float], + seed: int | None = None, + get_index: bool = False, ): """Splits the dataset by sample for conformal prediction. From e617ae3eeec8b37b5695127be295fda0e45569d5 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:58:05 +0530 Subject: [PATCH 09/13] Update __init__.py --- pyhealth/datasets/__init__.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index af55a227a..7387b5f68 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -1,3 +1,5 @@ + + class BaseEHRDataset: """This class is deprecated and should not be used.""" @@ -97,3 +99,19 @@ def __init__(self, *args, **kwargs): save_processors, ) from .collate import collate_temporal +from .splitter import ( + PatientSplit as PatientSplit, + Split as Split, + sample_balanced, + split_by_patient, + split_by_patient_conformal, + split_by_patient_conformal_tuh, + split_by_patient_tuh, + split_by_predefined, + split_by_sample, + split_by_sample_conformal, + split_by_sample_conformal_tuh, + split_by_sample_tuh, + split_by_visit, + split_by_visit_conformal, +) From 373ec9f094d1ba57ba12d1b4d8d3f9f3bf766beb Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:00:14 +0530 Subject: [PATCH 10/13] Add test for split_by_predefined function Added a test function to verify the callable nature of split_by_predefined. --- pyhealth/datasets/splitter.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 74f1240c3..1a91eeea3 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -666,6 +666,9 @@ def split_by_predefined( return tuple(torch.tensor(p) for p in parts) return tuple(dataset.subset(p) for p in parts) # type: ignore +def test_split_by_predefined_is_exported(): + from pyhealth.datasets import split_by_predefined + assert callable(split_by_predefined) def split_by_patient_conformal_tuh( dataset: SampleDataset, From 2b89b794965315fa0fb1315e13686ef8b93f5e23 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:00:58 +0530 Subject: [PATCH 11/13] Clean up CONTRIBUTING.md by removing old content Removed outdated guidelines and community information from the CONTRIBUTING.md file. --- CONTRIBUTING.md | 88 ------------------------------------------------- 1 file changed, 88 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 408224285..2c757d651 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -84,91 +84,3 @@ growing one near-identical function per dataset, port official splits as follows `train_value` form the held-out test set. - `by_patient=True` keeps all of a patient's samples in the same split (the dataset must have `patient_to_index`). Patients are sorted by ID before - shuffling, so the result does not depend on sample order. - - `split_key` (default `"split"`) and `train_value` (default `"train"`) adapt - the function to a different field name or partition label. - - `get_index=True` returns `torch.Tensor` index vectors instead of datasets. -3. **Wrap it only if the arguments are fixed.** If a dataset always needs the same - arguments, add a thin wrapper next to that dataset (in its module under - `pyhealth/datasets/`), not in `splitter.py`. The existing `*_tuh` functions in - `splitter.py` are such wrappers, kept for backward compatibility. -4. **Test it.** Add a unit test that checks the held-out set is exactly the - official eval partition, that no patient appears in more than one part when - `by_patient=True`, and that the same seed gives the same split. - -```python ->>> from pyhealth.datasets import split_by_predefined ->>> train, val, test = split_by_predefined( # doctest: +SKIP -... sample_dataset, [0.8, 0.2], seed=42, by_patient=True -... ) ->>> train, val, cal, test = split_by_predefined( # doctest: +SKIP -... sample_dataset, [0.6, 0.2, 0.2], seed=42, by_patient=True -... ) -``` - -These functions split an already-processed `SampleDataset`, so its processors -(vocabularies, learned statistics) were fitted on all samples. If preprocessing -must be fitted on the training part only, split inside `set_task` instead, with -`PatientSplit` or a custom `Split`. - -## Community - -We welcome you to join our community -via [Slack](https://join.slack.com/t/pyhealthworkspace/shared_invite/zt-1j8h4da47-ZZWENUYax7SlgCRHNyL2DQ) -workspace, -where you can ask questions, discuss ideas, and share your work. - -## Core Development - -We use [GitHub Flow](https://docs.github.com/en/get-started/quickstart/github-flow) -for development. - -### Branches - -We keep the following branches for parallel development and release management. - -- Main branches - - `master`: always reflects the production-ready state. - - `develop`: reflects the latest delivered development state. When the source code - in the `develop` branch reaches a stable point and is ready to be released, all - of the changes will be merged back into `master`. -- Release branches - - `*-release`: holds the code for a specific release. These branches are - created from `develop`, and will be merged back into `master` and `develop` - when the release is deployed. Only minor bug fixes or preparing meta-data for - a release (version number, build dates, etc.) are allowed in this branch. - - `archived`: holds the original version of PyHealth developed - by [Yue Zhao](https://www.andrew.cmu.edu/user/yuezhao2/) - et al. This branch is not maintained anymore. -- Feature branches: - - `[feature_name]`: feature branches are used to develop new features for the - upcoming release. These branches are created from `develop` and will be merged - back into `develop` once the feature is completed. -- Hotfix branches: - - `[hotfix_name]`: hotfix branches are used to quickly patch small bugs in the - production version. These branches are created from `master` and will be merged - back into `master` and `develop`. - -See this blog post for more details -on [Git Branching Model](https://nvie.com/posts/a-successful-git-branching-model/). - -### Pull Requests - -All code changes happen through pull requests. - -1. Create a new branch from `develop` (for feature) or `master` (for hotfix) -2. Make the changes -3. Write documentation for the changes -4. Write unit tests for the changes -5. Create a pull request to `develop` (for feature) or `master` branch (for hotfix) -6. Wait for the review and merge - -### Code Review - -We enforce code reviews for the main branches. - -- `master`: requires two approvals from the core developers -- `develop`: requires one approval from the core developers -Join our -[Slack workspace](https://join.slack.com/t/pyhealthworkspace/shared_invite/zt-1j8h4da47-ZZWENUYax7SlgCRHNyL2DQ) -to ask questions, discuss ideas, and share your work. From 19b3dda0d8fd9dbed68204442f0738896033b251 Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:23:52 +0530 Subject: [PATCH 12/13] Update temple_university_EEG_tasks.py --- pyhealth/tasks/temple_university_EEG_tasks.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/pyhealth/tasks/temple_university_EEG_tasks.py b/pyhealth/tasks/temple_university_EEG_tasks.py index 122f99311..fe8673c7d 100644 --- a/pyhealth/tasks/temple_university_EEG_tasks.py +++ b/pyhealth/tasks/temple_university_EEG_tasks.py @@ -20,7 +20,8 @@ class EEGEventsTUEV(BaseTask): - "signal": np.ndarray, shape (16, 256*5) - "offending_channel": int - "label": int - - "split": str, "train" or "eval", the dataset's official partition. Used by pyhealth.datasets.split_by_predefined. + - "split": str, "train" or "eval", the dataset's official partition. + Used by pyhealth.datasets.split_by_predefined. Examples: >>> from pyhealth.datasets import TUEVDataset @@ -148,7 +149,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: Iterates over both 'train' and 'eval' splits. Each sample includes a 'split' field and a precomputed 'stft' tensor. - Expected patient events to include a `signal_file` attribute pointing to an .edf file. + Expected patient events to include a `signal_file` attribute pointing to + an .edf file. """ pid = patient.patient_id samples: List[Dict[str, Any]] = [] @@ -199,8 +201,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: samples.append(sample) return samples - - + class EEGAbnormalTUAB(BaseTask): """Binary classification task for abnormal EEG detection on TUAB. @@ -275,12 +276,12 @@ def read_and_process_edf(fileName: str, @staticmethod def convert_to_bipolar(raw_data: np.ndarray, ch_name: List[str]) -> np.ndarray: """Convert raw EEG channels to 16 bipolar montage channels following BIOT. - + Returns: np.ndarray of shape (16, n_samples) """ channeled_data = np.zeros((16, raw_data.shape[1])) - + channeled_data[0] = ( raw_data[ch_name.index("EEG FP1-REF")] - raw_data[ch_name.index("EEG F7-REF")] @@ -351,8 +352,9 @@ def convert_to_bipolar(raw_data: np.ndarray, ch_name: List[str]) -> np.ndarray: def __call__(self, patient: Any) -> List[Dict[str, Any]]: """Processes one patient. Creates one 10-second window sample per segment. - Iterates over both 'train' and 'eval' splits. - Each sample includes a 'split' field and, when compute_stft=True, a precomputed 'stft' tensor. + Iterates over both 'train' and 'eval' splits. + Each sample includes a 'split' field and, when compute_stft=True, a + precomputed 'stft' tensor. """ pid = patient.patient_id samples: List[Dict[str, Any]] = [] @@ -413,7 +415,3 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) return samples - - - - From 6a2146590e3e44b6a03ad08fb8484c9dd5ff390d Mon Sep 17 00:00:00 2001 From: Aprameyaraya <91876432+Aprameyaraya@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:26:17 +0530 Subject: [PATCH 13/13] Update splitter.py --- pyhealth/datasets/splitter.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index b587e71f7..add8f57ef 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -792,9 +792,6 @@ def split_by_predefined( return tuple(torch.tensor(p) for p in parts) return tuple(dataset.subset(p) for p in parts) # type: ignore -def test_split_by_predefined_is_exported(): - from pyhealth.datasets import split_by_predefined - assert callable(split_by_predefined) def split_by_patient_conformal_tuh( dataset: SampleDataset,