pyhealth.datasets.MEDSDataset#

Dataset class for data in the Medical Event Data Standard (MEDS), a minimal event-based schema for machine learning over EHR data (MEDS Working Group / Arnrich et al., ICLR 2024 Workshop on Learning from Time Series For Health; openreview:IsHy2ebjIG). Sharded Parquet event files are read with their native types, and standard MEDS splits (train / tuning / held_out) can be selected directly via the subset argument. The canonical subject-to-split mapping is defined in metadata/subject_splits.parquet; see the MEDS schema documentation.

class pyhealth.datasets.MEDSDataset(root, tables=None, subset='all', split_source='metadata', dataset_name=None, config_path=None, **kwargs)[source]#

Bases: BaseDataset

Dataset for MEDS (Medical Event Data Standard) sources.

MEDS data is distributed as sharded, typed Parquet under per-split directories (data/train/*.parquet, data/tuning/*.parquet, data/held_out/*.parquet) plus a canonical subject-to-split map at metadata/subject_splits.parquet. See the MEDS schema documentation: https://medical-event-data-standard.github.io/

time must be a timezone-naive timestamp (MEDS reference schema); violations raise TypeError at construction.

Split handling:

The canonical split is available two ways, both optional:

as events: load the subject_splits table (tables=["meds", "subject_splits"]) and each subject carries one subject_splits event with attribute subject_splits/split – the exact pattern of EHRShot’s splits table, usable from Task.pre_filter or per-patient logic;
as a loader filter: subset="train" (or "tuning" / "held_out") keeps only that split’s patients in every loaded table, via the same patient-isin mechanic as dev mode.

split_source controls where subset gets its patient list: "metadata" (default) reads the canonical mapping file – authoritative per the MEDS spec and independent of directory layout; "directory" derives it from which data/<split>/ directory subjects appear in – useful when an export omits the metadata file. The two sources should agree; whether PyHealth must verify that equivalence is an open question for the upstream maintainer, so this class does not silently pick one when they could diverge: it uses exactly the source you asked for, and caches them separately.

Note

event_type is the table name ("meds") for every row; the clinically meaningful event kind lives in the meds/code attribute. This mirrors EHRShot, whose single ehrshot table also carries an event vocabulary in a code attribute. Whether upstream prefers mapping MEDS code onto event_type instead is a design question for the maintainer.

Parameters:

root (str) – Root directory of the MEDS dataset (the directory that contains data/ and metadata/).
tables (Optional[list[str]]) – Tables to load, as named in configs/meds.yaml. Defaults to ["meds"]; add "subject_splits" to expose the canonical split as events.
subset (str) – "train", "tuning", "held_out", or "all" (default). Anything but "all" filters every loaded table to that split’s patients.
split_source (Literal[‘metadata’, ‘directory’]) – Where subset gets its patient list from; see above. Ignored when subset="all".
dataset_name (Optional[str]) – Dataset name. Defaults to "meds".
config_path (Optional[str]) – Path to the YAML config. Defaults to configs/meds.yaml.
**kwargs – Forwarded to BaseDataset (cache_dir, num_workers, dev). Note dev mode’s 1000-patient cap is applied downstream of load_table (in BaseDataset._event_transform), so it composes with subset with no extra handling here.

Examples

>>> from pyhealth.datasets import MEDSDataset
>>> dataset = MEDSDataset(
...     root="/path/to/mimic-iv-demo-meds/0.0.1",
... )  
>>> dataset.stats()  
>>> # Canonical training split only, split map exposed as events:
>>> train = MEDSDataset(
...     root="/path/to/mimic-iv-demo-meds/0.0.1",
...     tables=["meds", "subject_splits"],
...     subset="train",
... )  

load_data()[source]#

Load all configured tables, restricted to the subset if any.

Returns:: The concatenated event frame, filtered to the subjects of self.subset when a split was requested.
Return type:: dd.DataFrame

clean_tmpdir()#

Cleans up the temporary directory within the cache.

Return type:: None

create_tmpdir()#

Creates and returns a new temporary directory within the cache.

Returns:: The path to the new temporary directory.
Return type:: Path

property default_task: Optional[BaseTask]#

Returns the default task for the dataset.

Returns:: The default task, if any.
Return type:: Optional[BaseTask]

get_patient(patient_id)#

Retrieves a Patient object for the given patient ID.

Parameters:: patient_id (str) – The ID of the patient to retrieve.
Returns:: The Patient object for the given ID.
Return type:: Patient
Raises:: AssertionError – If the patient ID is not found in the dataset.

property global_event_df: LazyFrame#

Returns the path to the cached event dataframe.

Returns:: The path to the cached event dataframe.
Return type:: Path

iter_patients(df=None)#

Yields Patient objects for each unique patient in the dataset.

Yields:: Iterator[Patient] – An iterator over Patient objects.
Return type:: Iterator[Patient]

load_table(table_name)#

Loads a table and processes joins if specified.

Parameters:

table_name (str) – The name of the table to load.

Returns:

The processed Dask dataframe for the table.

Return type:

dd.DataFrame

Raises:

ValueError – If the table is not found in the config.
FileNotFoundError – If the source file (CSV/TSV or Parquet) for the table or join is not found.

set_task(task=None, num_workers=None, input_processors=None, output_processors=None)#

Processes the base dataset to generate the task-specific sample dataset. The cache structure is as follows:

{task_name}_{task_uuid}/        # Cached data for specific task based on task name, schema, and args
    task_df.ld/                 # Intermediate task dataframe based on schema
    samples_{proc_uuid}.ld/     # Final processed samples after applying processors
        schema.pkl              # Saved SampleBuilder schema
        *.bin                   # Processed sample files

Parameters:

task (Optional[BaseTask]) – The task to set. Uses default task if None.
num_workers (int) – Number of workers for multi-threading. Default is self.num_workers.
input_processors (Optional[Dict[str, FeatureProcessor]]) – Pre-fitted input processors. If provided, these will be used instead of creating new ones from task’s input_schema. Defaults to None.
output_processors (Optional[Dict[str, FeatureProcessor]]) – Pre-fitted output processors. If provided, these will be used instead of creating new ones from task’s output_schema. Defaults to None.

Returns:

The generated sample dataset.

Return type:

SampleDataset

Raises:

AssertionError – If no default task is found and task is None.

stats()#

Prints statistics about the dataset.

Return type:: None

property unique_patient_ids: List[str]#

Returns a list of unique patient IDs.

Returns:: List of unique patient IDs.
Return type:: List[str]