pyhealth.datasets.splitter#

Several data splitting function for pyhealth.datasets module to obtain training / validation / test sets.

The patient-level functions (split_by_patient, split_by_patient_conformal and their TUH variants) sort patient IDs before the seeded shuffle, so a given seed selects the same patients however the samples are ordered. Releases before this change shuffled patients in sample order, so the same seed can now select different patients than it did in those releases.

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:

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 <unk>, 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. 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 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.

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:

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"]))

class pyhealth.datasets.splitter.ABC[source]#

Bases: object

Helper class that provides a standard way to create an ABC using inheritance.

pyhealth.datasets.splitter.abstractmethod(funcobj)[source]#

A decorator indicating abstract methods.

Requires that the metaclass is ABCMeta or derived from it. A class that has a metaclass derived from ABCMeta cannot be instantiated unless all of its abstract methods are overridden. The abstract methods can be called using any of the normal ‘super’ call mechanisms. abstractmethod() may be used to declare abstract methods for properties and descriptors.

Usage:

class C(metaclass=ABCMeta):
@abstractmethod def my_abstract_method(self, arg1, arg2, argN):

…

class pyhealth.datasets.splitter.Mapping#

Bases: Collection

A Mapping is a generic container for associating key/value pairs.

This class provides concrete generic implementations of all methods except for __getitem__, __iter__, and __len__.

get(k[, d]) → D[k] if k in D, else d. d defaults to None.#

items() → a set-like object providing a view on D's items#

keys() → a set-like object providing a view on D's keys#

values() → an object providing a view on D's values#

class pyhealth.datasets.splitter.Sequence#

Bases: Reversible, Collection

All the operations on a read-only sequence.

Concrete subclasses must override __new__ or __init__, __getitem__, and __len__.

count(value) → integer -- return number of occurrences of value#

index(value[, start[, stop]]) → integer -- return first index of value.#

Raises ValueError if the value is not present.

Supporting start and stop arguments is optional, but recommended.

pyhealth.datasets.splitter.dataclass(cls=None, /, *, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False)[source]#

Add dunder methods based on the fields defined in the class.

Examines PEP 526 __annotations__ to determine fields.

If init is true, an __init__() method is added to the class. If repr is true, a __repr__() method is added. If order is true, rich comparison dunder methods are added. If unsafe_hash is true, a __hash__() method is added. If frozen is true, fields may not be assigned to after instance creation. If match_args is true, the __match_args__ tuple is added. If kw_only is true, then by default all fields are keyword-only. If slots is true, a new class with a __slots__ attribute is returned.

class pyhealth.datasets.splitter.chain#

Bases: object

chain(*iterables) –> chain object

Return a chain object whose .__next__() method returns elements from the first iterable until it is exhausted, then elements from the next iterable, until all of the iterables are exhausted.

from_iterable()#: Alternative chain() constructor taking a single iterable argument that evaluates lazily.

class pyhealth.datasets.splitter.SampleDataset(path, dataset_name=None, task_name=None, **kwargs)[source]#

Bases: StreamingDataset

A streaming dataset that loads sample metadata and processors from disk.

SampleDataset expects the path directory to contain a schema.pkl file created by a SampleBuilder.save(…) call. The schema.pkl must include the fitted input_schema, output_schema, input_processors, output_processors, patient_to_index and record_to_index mappings.

input_schema#: The configuration used to instantiate processors for input features (string aliases or processor specs).

output_schema#: The configuration used to instantiate processors for output features.

input_processors#: A mapping of input feature names to fitted FeatureProcessor instances.

output_processors#: A mapping of output feature names to fitted FeatureProcessor instances.

patient_to_index#: Dictionary mapping patient IDs to the list of sample indices associated with that patient.

record_to_index#: Dictionary mapping record/visit IDs to the list of sample indices associated with that record.

dataset_name#: Optional human friendly dataset name.

task_name#: Optional human friendly task name.

Examples

>>> from pyhealth.datasets import create_sample_dataset
>>> dataset = create_sample_dataset(  
...     samples=[
...         {"patient_id": "p1", "feature": 1, "label": 0},
...         {"patient_id": "p2", "feature": 2, "label": 1},
...     ],
...     input_schema={"feature": "raw"},
...     output_schema={"label": "raw"},
...     in_memory=False,
... )
>>> dataset.subset([1]).patient_to_index  
{'p2': [0]}
>>> dataset.close()  

subset(indices)[source]#

Create a StreamingDataset restricted to the provided indices.

Return type:: SampleDataset

close()[source]#

Cleans up any temporary directories used by the dataset.

Return type:: None

get_len(num_workers, batch_size)#

Return type:: int

load_state_dict(state_dict)#

Return type:: None

property on_demand_bytes: bool#

Return type:: bool

reset()#

Return type:: None

reset_state_dict()#

Return type:: None

set_batch_size(batch_size)#

Return type:: None

set_drop_last(drop_last)#

Set the drop_last parameter.

Invalidates the shuffler cache when the parameter changes to ensure subsequent length calculations reflect the new drop_last setting.

Parameters:: drop_last (bool) – Whether to drop the last incomplete batch.
Return type:: None

set_epoch(current_epoch)#

Set the current epoch to the dataset on epoch starts.

When using the StreamingDataLoader, this is done automatically

Return type:: None

set_num_workers(num_workers)#

Return type:: None

set_shuffle(shuffle)#

Set the shuffle parameter.

Invalidates the shuffler cache when the parameter changes to ensure subsequent length calculations reflect the new shuffle setting.

Parameters:: shuffle (bool) – Whether to shuffle the dataset.
Return type:: None

state_dict(num_samples_yielded, num_workers, batch_size)#

Return type:: dict[str, Any]

class pyhealth.datasets.splitter.Split[source]#

Bases: 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. 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])]

abstract split_indices(patient_to_index)[source]#

Returns the sample indices of each part.

Parameters:: patient_to_index (Mapping[str, Sequence[int]]) – Patient ID to the indices of its samples.
Return type:: list[ndarray]
Returns:: One int64 array of sample indices per part, training part first. Sorted indices keep reads from the cache sequential.

abstract to_dict()[source]#

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.

Return type:: dict

class pyhealth.datasets.splitter.PatientSplit(ratios, seed=0)[source]#

Bases: Split

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.

Unlike 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.

Parameters:

ratios (Sequence[float]) – 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 (int) – 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#

to_dict()[source]#

JSON-serialisable description, used in cache keys and metadata.

Return type:: dict

split_indices(patient_to_index)[source]#

Splits sample indices by patient.

Patients are sorted, shuffled with seed and cut by ratios.

Parameters:: patient_to_index (Mapping[str, Sequence[int]]) – Patient ID to the indices of its samples.
Return type:: list[ndarray]
Returns:: One sorted int64 array of sample indices per ratio, in order (train first). Each patient’s samples land in exactly one part.

pyhealth.datasets.splitter.sample_balanced(dataset, ratio=1.0, subsample=1.0, seed=None)[source]#

Keep positives and negatives at a target ratio, then cap total size.

Parameters:

dataset (SampleDataset) – Dataset with patient_to_index populated.
ratio (float) – Negatives per positive (e.g., 1.0 -> ~1 neg per pos). Values <=0 keep only positives.
subsample (float) – Max fraction of the original dataset size to retain. If the ratio-selected set exceeds len(dataset) * subsample, both positives and negatives are downsampled proportionally while preserving the ratio as closely as possible.
seed (Optional[int]) – Optional RNG seed for reproducible negative sampling.

Return type:

SampleDataset

Returns:

A new SampleDataset containing all positives plus sampled negatives, with refreshed patient_to_index and record_to_index mappings.

pyhealth.datasets.splitter.split_by_visit(dataset, ratios, seed=None)[source]#

Splits the dataset by visit (i.e., samples).

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float], List[float]]) – a list/tuple of ratios for train / val / test
seed (Optional[int]) – random seed for shuffling the dataset

Returns:

SampleDataset views: (dataset.subset) over the same processed samples.

Return type:

train_dataset, val_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_visit
>>> train, val, test = split_by_visit(  
...     samples, [0.8, 0.1, 0.1], 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 PatientSplit or a custom Split.

pyhealth.datasets.splitter.split_by_patient(dataset, ratios, seed=None)[source]#

Splits the dataset by patient.

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float], List[float]]) – a list/tuple of ratios for train / val / test
seed (Optional[int]) – random seed for shuffling the dataset. Patients are sorted by ID before shuffling, so the split depends only on the seed and the set of patients, not on the order of the samples.

Returns:

SampleDataset views: (dataset.subset) over the same processed samples.

Return type:

train_dataset, val_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_patient
>>> train, val, test = split_by_patient(  
...     samples, [0.8, 0.1, 0.1], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_sample(dataset, ratios, seed=None, get_index=False)[source]#

Splits the dataset by sample

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float], List[float]]) – a list/tuple of ratios for train / val / test
seed (Optional[int]) – random seed for shuffling the dataset
get_index (Optional[bool]) – if True, return three torch.Tensor index vectors instead of datasets

Returns:

SampleDataset views: (dataset.subset) over the same processed samples, or index tensors if get_index is True.

Return type:

train_dataset, val_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_sample
>>> train, val, test = split_by_sample(  
...     samples, [0.8, 0.1, 0.1], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_visit_conformal(dataset, ratios, seed=None)[source]#

Splits the dataset by visit (i.e., samples) for conformal prediction.

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float, float], List[float]]) – a list/tuple of ratios for train / val / cal / test
seed (Optional[int]) – random seed for shuffling the dataset

Returns:

SampleDataset views (dataset.subset) over the same processed samples.

Return type:

train_dataset, val_dataset, cal_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_visit_conformal
>>> train, val, cal, test = split_by_visit_conformal(  
...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_patient_conformal(dataset, ratios, seed=None)[source]#

Splits the dataset by patient for conformal prediction.

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float, float], List[float]]) – a list/tuple of ratios for train / val / cal / test
seed (Optional[int]) – random seed for shuffling the dataset

Returns:

SampleDataset views (dataset.subset) over the same processed samples.

Return type:

train_dataset, val_dataset, cal_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_patient_conformal
>>> train, val, cal, test = split_by_patient_conformal(  
...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_patient_conformal_tuh(dataset, ratios, seed=None, get_index=False)[source]#

Splits a TUH EEG dataset by PATIENT within the TUH train partition.

Unlike split_by_sample_conformal_tuh(), which shuffles windows independently, this function ensures that all windows from a given patient stay in the same split (train, val, or cal). This prevents within-patient data leakage between training and calibration — a critical requirement for honest evaluation of conformal prediction methods, especially NCP, which relies on nearest-neighbor search in embedding space.

The TUH corpus is designed so that each patient appears exclusively in either the training partition (265 patients for TUSZ) or the evaluation partition (50 patients), never both. The test set is therefore always the TUH eval partition (fixed, fully held-out patients).

Parameters:

dataset (SampleDataset) – a SampleDataset produced by EEGEventsTUEV or EEGAbnormalTUAB. Each sample must carry a "split" field ("train" or "eval") and the dataset must have a patient_to_index mapping.
ratios (Union[Tuple[float, float, float], List[float]]) – fraction of patients in the TUH train partition to assign to train / val / cal respectively. Must be a length-3 sequence summing to 1.0.
seed (Optional[int]) – random seed used to shuffle the patient list.
get_index (Optional[bool]) – if True, return four torch.Tensor index vectors 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(  
...     samples, [0.6, 0.2, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_sample_conformal_tuh(dataset, ratios, seed=None, get_index=False)[source]#

Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined train/eval split.

Parameters:

dataset (SampleDataset) – a SampleDataset object produced by EEGEventsTUEV or EEGAbnormalTUAB
ratios (Union[Tuple[float, float, float], List[float]]) – the fraction of the train pool assigned to train / val / cal respectively
seed (Optional[int]) – random seed for shuffling the train pool
get_index (Optional[bool]) – if True, return four torch.Tensor index vectors instead 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(  
...     samples, [0.6, 0.2, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_patient_tuh(dataset, ratios, seed=None, get_index=False)[source]#

Splits a TUH EEG dataset by PATIENT within the TUH train partition.

Like split_by_patient_conformal_tuh() but returns a 3-way split (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.

Ensures that all windows from a given patient stay in the same split (train or val), preventing within-patient data leakage.

Parameters:

dataset (SampleDataset) – a SampleDataset produced by EEGEventsTUEV or EEGAbnormalTUAB. Each sample must carry a "split" field ("train" or "eval") and the dataset must have a patient_to_index mapping.
ratios (Union[Tuple[float, float], List[float]]) – fraction of patients in the TUH train partition to assign to train / val respectively. Must be a length-2 sequence summing to 1.0.
seed (Optional[int]) – random seed used to shuffle the patient list.
get_index (Optional[bool]) – if True, return three torch.Tensor index 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(  
...     samples, [0.8, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_sample_tuh(dataset, ratios, seed=None, get_index=False)[source]#

Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined split.

Like split_by_sample_conformal_tuh() but returns a 3-way split (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.

Parameters:

dataset (SampleDataset) – a SampleDataset object produced by EEGEventsTUEV or EEGAbnormalTUAB
ratios (Union[Tuple[float, float], List[float]]) – the fraction of the train pool assigned to train / val respectively. Must be a length-2 sequence summing to 1.0.
seed (Optional[int]) – random seed for shuffling the train pool
get_index (Optional[bool]) – if True, return three torch.Tensor index vectors instead 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(  
...     samples, [0.8, 0.2], seed=42
... )

Note

pyhealth.datasets.splitter.split_by_sample_conformal(dataset, ratios, seed=None, get_index=False)[source]#

Splits the dataset by sample for conformal prediction.

Parameters:

dataset (SampleDataset) – a SampleDataset object
ratios (Union[Tuple[float, float, float, float], List[float]]) – a list/tuple of ratios for train / val / cal / test
seed (Optional[int]) – random seed for shuffling the dataset
get_index (Optional[bool]) – if True, return indices instead of datasets

Returns:

SampleDataset views (dataset.subset) over the same processed samples, or four index tensors if get_index is True.

Return type:

train_dataset, val_dataset, cal_dataset, test_dataset

Examples

>>> from pyhealth.datasets import split_by_sample_conformal
>>> train, val, cal, test = split_by_sample_conformal(  
...     samples, [0.6, 0.1, 0.1, 0.2], seed=42
... )

Note