diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index c9a88b7ff..007014a5b 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -79,6 +79,9 @@ future runs without re-fitting. - ``samples.record_to_index`` — maps a visit/record ID to the sample indices for that visit. +Calling ``samples.subset(...)`` rebuilds both lookups with indices local to the +new dataset, so they remain valid after repeated splitting. + For testing or small cohorts you can skip the disk step entirely using ``InMemorySampleDataset``, which holds all processed samples in RAM and is returned by default from ``create_sample_dataset()``. diff --git a/examples/sample_dataset_subset.py b/examples/sample_dataset_subset.py new file mode 100644 index 000000000..3342f7b1f --- /dev/null +++ b/examples/sample_dataset_subset.py @@ -0,0 +1,27 @@ +"""Create reordered and nested SampleDataset subsets.""" + +from pyhealth.datasets import create_sample_dataset + + +def main(): + samples = [ + {"patient_id": "p1", "record_id": "r1", "feature": 1, "label": 0}, + {"patient_id": "p2", "record_id": "r2", "feature": 2, "label": 1}, + {"patient_id": "p1", "record_id": "r3", "feature": 3, "label": 0}, + {"patient_id": "p3", "record_id": "r4", "feature": 4, "label": 1}, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"feature": "raw"}, + output_schema={"label": "raw"}, + ) + + subset = dataset.subset([3, 0, 2]) + nested_subset = subset.subset([2, 0]) + + print(subset.patient_to_index) + print(nested_subset.patient_to_index) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 134bede35..6b9a751ee 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -15,6 +15,22 @@ from ..processors.base_processor import FeatureProcessor +def _remap_index_mapping( + mapping: dict[str, list[int]], indices: Sequence[int] +) -> dict[str, list[int]]: + key_by_index = { + index: key + for key, mapped_indices in mapping.items() + for index in mapped_indices + } + remapped: dict[str, list[int]] = {} + for new_index, old_index in enumerate(indices): + key = key_by_index.get(old_index) + if key is not None: + remapped.setdefault(key, []).append(new_index) + return remapped + + class SampleBuilder: """Fit feature processors and transform pickled samples without materializing a dataset. @@ -275,6 +291,21 @@ class SampleDataset(litdata.StreamingDataset): 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( # doctest: +SKIP + ... 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 # doctest: +SKIP + {'p2': [0]} + >>> dataset.close() # doctest: +SKIP """ def __init__( @@ -355,6 +386,8 @@ def subset(self, indices: Union[Sequence[int], slice]) -> "SampleDataset": if isinstance(indices, slice): indices = range(*indices.indices(dataset_length)) + indices = list(indices) + if any(idx < 0 or idx >= dataset_length for idx in indices): raise ValueError( f"Subset indices must be in [0, {dataset_length - 1}] for the provided dataset." @@ -403,6 +436,12 @@ def subset(self, indices: Union[Sequence[int], slice]) -> "SampleDataset": new_dataset.subsampled_files = new_subsampled_files new_dataset.region_of_interest = new_roi + new_dataset.patient_to_index = _remap_index_mapping( + self.patient_to_index, indices + ) + new_dataset.record_to_index = _remap_index_mapping( + self.record_to_index, indices + ) new_dataset.reset() return new_dataset @@ -430,6 +469,16 @@ class InMemorySampleDataset(SampleDataset): for fast, repeated access to samples without disk I/O, at the cost of higher memory usage. + Examples: + >>> from pyhealth.datasets import create_sample_dataset + >>> dataset = create_sample_dataset( + ... samples=[{"patient_id": "p1", "feature": 1, "label": 0}], + ... input_schema={"feature": "raw"}, + ... output_schema={"label": "raw"}, + ... ) + >>> dataset[0]["patient_id"] + 'p1' + Note: This class is intended for testing and debugging purposes where dataset sizes are small enough to fit into memory. @@ -522,12 +571,23 @@ def __iter__(self) -> Iterable[Dict[str, Any]]: # type: ignore def subset(self, indices: Union[Sequence[int], slice]) -> SampleDataset: if isinstance(indices, slice): + subset_indices = list(range(*indices.indices(len(self)))) samples = self._data[indices] else: - samples = [self._data[i] for i in indices] + raw_indices = list(indices) + samples = [self._data[i] for i in raw_indices] + subset_indices = [ + index if index >= 0 else len(self) + index for index in raw_indices + ] new_dataset = copy.deepcopy(self) new_dataset._data = samples + new_dataset.patient_to_index = _remap_index_mapping( + self.patient_to_index, subset_indices + ) + new_dataset.record_to_index = _remap_index_mapping( + self.record_to_index, subset_indices + ) return new_dataset def close(self) -> None: diff --git a/tests/core/test_sample_dataset.py b/tests/core/test_sample_dataset.py index 8dd4629c0..95986bde3 100644 --- a/tests/core/test_sample_dataset.py +++ b/tests/core/test_sample_dataset.py @@ -83,6 +83,30 @@ def test_subset_slice(self): for d, m in zip(list_disk, list_mem): self.assertEqual(d["feature"], m["feature"]) + def test_subset_rebuilds_index_mappings(self): + ds_disk, ds_mem = self._get_datasets() + + for dataset in (ds_disk, ds_mem): + subset = dataset.subset([5, 2, 15]) + self.assertEqual( + subset.patient_to_index, + {"p5": [0], "p2": [1], "p15": [2]}, + ) + self.assertEqual( + subset.record_to_index, + {"r5": [0], "r2": [1], "r15": [2]}, + ) + + nested_subset = subset.subset([2, 0]) + self.assertEqual( + nested_subset.patient_to_index, + {"p15": [0], "p5": [1]}, + ) + self.assertEqual( + nested_subset.record_to_index, + {"r15": [0], "r5": [1]}, + ) + def test_set_shuffle(self): ds_disk, ds_mem = self._get_datasets()