pyhealth.tasks.generate_ehr

Tasks that turn a longitudinal EHR dataset into training samples for unconditional synthetic-EHR generators, plus helpers to flatten generated output into the long-form dataframe consumed by pyhealth.metrics.generative.

The classes are flat and independent – one per (model family, dataset). The extraction is MIMIC-shaped, assuming an admissions event type and a hadm_id linking codes to an admission, so the dataset is named in the class and a task for eICU/OMOP/MEDS belongs alongside these rather than below them.

Match the task to the model: each generator family reads its codes in a different shape, and handing a model the wrong shape fails silently rather than loudly.

Task	Encoding	Models
EHRGenerationMIMIC3 / MIMIC4	one multi-hot row per visit	HALO
EHRSequenceGenerationMIMIC3 / MIMIC4	per-visit code indices	GPT2, PromptEHR
EHRCodeSetGenerationMIMIC3 / MIMIC4	one code set per patient	MedGAN, CorGAN

Task Classes

class pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3(code_mapping=None)

Bases: BaseTask

Per-visit ICD-9 code sets from MIMIC-III as multi-hot rows. For HALO.

HALO’s transformer consumes a multi-hot vector per context position, so this hands it exactly that and nothing is repacked on the way in.

Patients with fewer than min_visits coded admissions are skipped.

Examples

>>> from pyhealth.datasets import MIMIC3Dataset
>>> from pyhealth.tasks import EHRGenerationMIMIC3
>>> dataset = MIMIC3Dataset(
...     root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"]
... )
>>> samples = dataset.set_task(EHRGenerationMIMIC3())
>>> samples[0]["visits"].shape  # (num_visits, vocab_size)
torch.Size([3, 512])

task_name: str = 'ehr_generation_mimic3'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.nested_multihot_processor.NestedMultiHotProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd9_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

class pyhealth.tasks.generate_ehr.EHRGenerationMIMIC4(code_mapping=None)

Bases: BaseTask

Per-visit ICD code sets from MIMIC-IV as multi-hot rows. For HALO.

MIMIC-IV’s diagnosis codes live on icd_code rather than MIMIC-III’s icd9_code; otherwise identical to EHRGenerationMIMIC3.

Examples

>>> from pyhealth.datasets import MIMIC4Dataset
>>> from pyhealth.tasks import EHRGenerationMIMIC4
>>> dataset = MIMIC4Dataset(
...     ehr_root="/path/to/mimiciv/2.2/",
...     ehr_tables=["patients", "admissions", "diagnoses_icd"],
... )
>>> samples = dataset.set_task(EHRGenerationMIMIC4())
>>> samples[0]["visits"].shape  # (num_visits, vocab_size)
torch.Size([3, 512])

task_name: str = 'ehr_generation_mimic4'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.nested_multihot_processor.NestedMultiHotProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

class pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC3(code_mapping=None)

Bases: BaseTask

Per-visit ICD-9 code indices from MIMIC-III. For GPT2 and PromptEHR.

Both are token-sequence models: they flatten each visit into a stream of code ids, so indices are what they want. Handing them the multi-hot form means encoding a code set and decoding it straight back.

Examples

>>> from pyhealth.datasets import MIMIC3Dataset
>>> from pyhealth.tasks import EHRSequenceGenerationMIMIC3
>>> dataset = MIMIC3Dataset(
...     root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"]
... )
>>> samples = dataset.set_task(EHRSequenceGenerationMIMIC3())
>>> samples[0]["visits"].shape  # (num_visits, max_codes_per_visit)
torch.Size([3, 12])

task_name: str = 'ehr_sequence_generation_mimic3'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.nested_sequence_processor.NestedSequenceProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd9_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

class pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC4(code_mapping=None)

Bases: BaseTask

Per-visit ICD code indices from MIMIC-IV. For GPT2 and PromptEHR.

Examples

>>> from pyhealth.datasets import MIMIC4Dataset
>>> from pyhealth.tasks import EHRSequenceGenerationMIMIC4
>>> dataset = MIMIC4Dataset(
...     ehr_root="/path/to/mimiciv/2.2/",
...     ehr_tables=["patients", "admissions", "diagnoses_icd"],
... )
>>> samples = dataset.set_task(EHRSequenceGenerationMIMIC4())
>>> samples[0]["visits"].shape  # (num_visits, max_codes_per_visit)
torch.Size([3, 12])

task_name: str = 'ehr_sequence_generation_mimic4'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.nested_sequence_processor.NestedSequenceProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

class pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC3(code_mapping=None)

Bases: BaseTask

One pooled ICD-9 code set per MIMIC-III patient. For MedGAN and CorGAN.

Bag-of-codes generators emit a single aggregate vector per patient, so the visit axis is collapsed here rather than inside the model. min_visits still counts real admissions, before the codes are pooled.

Note

With the visit axis gone, the next-visit utility metric in pyhealth.metrics.generative is not meaningful for these models; see this module’s header.

Examples

>>> from pyhealth.datasets import MIMIC3Dataset
>>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC3
>>> dataset = MIMIC3Dataset(
...     root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"]
... )
>>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC3())
>>> samples[0]["visits"].shape  # (vocab_size,)
torch.Size([512])

task_name: str = 'ehr_codeset_generation_mimic3'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.multi_hot_processor.MultiHotProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd9_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

class pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC4(code_mapping=None)

Bases: BaseTask

One pooled ICD code set per MIMIC-IV patient. For MedGAN and CorGAN.

Examples

>>> from pyhealth.datasets import MIMIC4Dataset
>>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC4
>>> dataset = MIMIC4Dataset(
...     ehr_root="/path/to/mimiciv/2.2/",
...     ehr_tables=["patients", "admissions", "diagnoses_icd"],
... )
>>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC4())
>>> samples[0]["visits"].shape  # (vocab_size,)
torch.Size([512])

task_name: str = 'ehr_codeset_generation_mimic4'

input_schema: ClassVar[dict[str, str | type]] = {'visits': <class 'pyhealth.processors.multi_hot_processor.MultiHotProcessor'>}

output_schema: ClassVar[dict[str, str | type]] = {}

event_type: str = 'diagnoses_icd'

code_attr: str = 'icd_code'

min_visits: int = 2

pre_filter(df)

Return type:

LazyFrame

Helper Functions

pyhealth.tasks.generate_ehr.decode_dataset(sample_dataset, feature_key='visits')

Decode a processed multi-hot SampleDataset back into code records.

Inverts the NestedMultiHotProcessor encoding using its vocabulary (skipping <pad>/<unk>), yielding one {"visits": [[code_str, ...], ...]} record per sample. Use this to build the real train/test frames that evaluate_synthetic_ehr compares against.

Codes come back in vocabulary order, not the order they were charted in, and repeats collapse – the multi-hot form records presence, not sequence or count within a visit.

Parameters:

• sample_dataset – A SampleDataset (or a split of one) produced by EHRGenerationMIMIC3 / EHRGenerationMIMIC4.

• feature_key (str) – Input feature key holding the nested code sequence. Default "visits".

Return type:

list[dict]

Returns:

List of {"visits": [[code_str, ...], ...]} records.

Raises:

TypeError – If feature_key is not backed by a NestedMultiHotProcessor.

Examples

>>> from pyhealth.tasks.generate_ehr import decode_dataset
>>> records = decode_dataset(samples)
>>> records[0]["visits"][0]
['4019', '25000']

pyhealth.tasks.generate_ehr.to_evaluation_dataframe(records, label_fn=None, subject_col='id', visit_col='time', code_col='visit_codes', label_col='labels')

Flatten EHR-generation records into the long-form evaluation dataframe.

Produces the one-row-per-(patient, visit, code) table consumed by pyhealth.metrics.generative.evaluate_synthetic_ehr() (and the utils.py / privacy.py / utility.py functions beneath it).

Subjects are numbered sequentially (0, 1, 2, …) in subject_col; any "patient_id" on the records is ignored, since synthetic patients do not correspond to real ones.

Parameters:

• records – Iterable of {"visits": [[code, ...], ...]} dicts. Both the generation tasks’ output and a generator’s generate() output have this shape.

• label_fn (Optional[Callable[dict, int]]) – Optional callable mapping a record to a binary patient label (0/1) used by the utility metrics. Defaults to all-zeros.

• subject_col (str) – Output patient-id column. Default "id".

• visit_col (str) – Output visit-index column. Default "time".

• code_col (str) – Output single-code column. Default "visit_codes".

• label_col (str) – Output binary-label column. Default "labels".

Returns:

pandas.DataFrame with columns [subject_col, visit_col, code_col, label_col].

Examples

>>> from pyhealth.tasks.generate_ehr import to_evaluation_dataframe
>>> records = [{"visits": [["4019", "25000"], ["4019"]]}]
>>> to_evaluation_dataframe(records)
   id  time visit_codes  labels
0   0     0        4019       0
1   0     0       25000       0
2   0     1        4019       0