diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b91c13a75..a553d5573 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -77,8 +77,23 @@ which runs `tools/check_pr_rules.py`. You can run the same check locally: python tools/check_pr_rules.py --base --head HEAD ``` -## Community - -Join our -[Slack workspace](https://join.slack.com/t/pyhealthworkspace/shared_invite/zt-1j8h4da47-ZZWENUYax7SlgCRHNyL2DQ) -to ask questions, discuss ideas, and share your work. +## 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 diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 5cc4cfc4f..659eb615d 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.""" @@ -99,3 +101,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, +) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 7b7013031..add8f57ef 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. @@ -507,9 +506,9 @@ def split_by_strat_fold( 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 @@ -570,8 +569,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. @@ -632,8 +631,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. @@ -701,11 +700,104 @@ def split_by_patient_conformal( return train_dataset, val_dataset, cal_dataset, test_dataset +def split_by_predefined( + dataset: SampleDataset, + 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", +): + """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]], - 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. @@ -721,6 +813,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 @@ -759,66 +853,19 @@ def split_by_patient_conformal_tuh( or a custom :class:`Split`. """ _validate_ratios(ratios, expected_length=3) - - # 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( 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 (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 @@ -852,58 +899,14 @@ def split_by_sample_conformal_tuh( or a custom :class:`Split`. """ _validate_ratios(ratios, expected_length=3) - - # 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( 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, ): """Splits a TUH EEG dataset by PATIENT within the TUH train partition. @@ -914,6 +917,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 @@ -951,58 +956,14 @@ def split_by_patient_tuh( or a custom :class:`Split`. """ _validate_ratios(ratios, expected_length=2) - - 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( 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, ): """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined split. @@ -1010,6 +971,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`` @@ -1044,51 +1007,14 @@ def split_by_sample_tuh( or a custom :class:`Split`. """ _validate_ratios(ratios, expected_length=2) - - 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( 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. diff --git a/pyhealth/tasks/temple_university_EEG_tasks.py b/pyhealth/tasks/temple_university_EEG_tasks.py index 13e8c7206..fe8673c7d 100644 --- a/pyhealth/tasks/temple_university_EEG_tasks.py +++ b/pyhealth/tasks/temple_university_EEG_tasks.py @@ -20,6 +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. Examples: >>> from pyhealth.datasets import TUEVDataset @@ -147,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]] = [] @@ -198,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. @@ -274,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")] @@ -350,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]] = [] @@ -412,7 +415,3 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) return samples - - - - 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)