diff --git a/docs/api/processors.rst b/docs/api/processors.rst index a06e3c955..e4262fa4c 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -26,6 +26,7 @@ Available Processors - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures) - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history) +- ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models) - ``NestedFloatsProcessor``: For nested numerical sequences with optional forward-fill **Label Processors:** @@ -283,6 +284,7 @@ Common string keys for automatic processor selection: - ``"temporal_timeseries"``: For time-series data with preserved timestamps (use in place of ``"timeseries"`` when building ``UnifiedMultimodalEmbeddingModel``) - ``"sequence"``: For categorical sequences (medical codes) - ``"nested_sequence"``: For nested categorical sequences (visit history) +- ``"nested_multihot"``: For nested categorical sequences as per-visit multi-hot vectors (set membership, not order) - ``"nested_sequence_floats"``: For nested numerical sequences - ``"binary"``: For binary labels - ``"multiclass"``: For multi-class labels @@ -477,6 +479,7 @@ API Reference processors/pyhealth.processors.DatasetProcessor processors/pyhealth.processors.SequenceProcessor processors/pyhealth.processors.NestedSequenceProcessor + processors/pyhealth.processors.NestedMultiHotProcessor processors/pyhealth.processors.NestedFloatsProcessor processors/pyhealth.processors.BinaryLabelProcessor processors/pyhealth.processors.MultiClassLabelProcessor diff --git a/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst b/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst new file mode 100644 index 000000000..00a743ae0 --- /dev/null +++ b/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst @@ -0,0 +1,21 @@ +pyhealth.processors.NestedMultiHotProcessor +=========================================== + +Processor for nested categorical sequences, emitted as per-visit multi-hot +vectors. + +Takes the same input as ``NestedSequenceProcessor`` -- a list of visits, each a +list of codes -- but emits one row per visit with one column per vocabulary +entry, set to 1 where the code is present. Repeats within a visit collapse, so +this records presence rather than order or count. + +Prefer it over ``NestedSequenceProcessor`` for set-membership models such as +generative EHR models: the index form pads every visit to the longest visit +seen during ``fit``, so a single outlier visit sets the width for the whole +dataset. Sizing by the vocabulary instead is both smaller and cheaper to +consume. + +.. autoclass:: pyhealth.processors.NestedMultiHotProcessor + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 735f6d39f..b250cd271 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -139,6 +139,9 @@ a quick reference: * - ``"nested_sequence"`` - ``NestedSequenceProcessor`` - Cumulative visit history (drug recommendation, readmission) + * - ``"nested_multihot"`` + - ``NestedMultiHotProcessor`` + - Per-visit code sets (HALO; see ``EHRGenerationMIMIC3``) * - ``"tensor"`` - ``TensorProcessor`` - Aggregated numeric values (e.g. last lab value per item) diff --git a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst index 77c332f33..796201c67 100644 --- a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst +++ b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst @@ -1,25 +1,66 @@ pyhealth.tasks.generate_ehr =========================================== -Task that turns a longitudinal EHR dataset into per-patient, per-visit code -sequences for training unconditional synthetic-EHR generators (HALO, GPT2, -PromptEHR, MedGAN, CorGAN), plus helpers to flatten generated output into the -long-form dataframe consumed by :mod:`pyhealth.metrics.generative`. +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 +:mod:`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. + +.. list-table:: + :header-rows: 1 + :widths: 38 32 30 + + * - 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 ------------ -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGeneration +.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3 +.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC4 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC4 +.. autoclass:: pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC3 + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC4 + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC3 + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC4 :members: :undoc-members: :show-inheritance: diff --git a/examples/gpt2_mimic3.py b/examples/gpt2_mimic3.py new file mode 100644 index 000000000..d70382030 --- /dev/null +++ b/examples/gpt2_mimic3.py @@ -0,0 +1,115 @@ +"""Example: train the GPT-2 baseline on MIMIC-III and generate patients. + +GPT2 and PromptEHR are token-sequence models: each patient becomes one flat +stream of code ids, ``[BOS] codes_v1 [DELIM] codes_v2 ... [EOS]``. They read +per-visit code *indices*, so this uses EHRSequenceGenerationMIMIC3 -- not the +multi-hot EHRGenerationMIMIC3 that HALO takes (see halo_mimic3.py). Pairing a +model with the wrong task does not raise; it trains on nonsense. + +Swap GPT2 for PromptEHR below and the rest of the script is unchanged. + +This example demonstrates: +1. Loading MIMIC-III data +2. Applying the EHRSequenceGenerationMIMIC3 task (per-visit ICD-9 code indices) +3. Training GPT2 with its custom training loop +4. Generating synthetic patients +5. Evaluating the synthetic data with the generative metrics suite +""" + +from pyhealth.datasets import MIMIC3Dataset, split_by_patient +from pyhealth.metrics.generative import evaluate_synthetic_ehr +from pyhealth.models import GPT2 +from pyhealth.tasks import EHRSequenceGenerationMIMIC3, to_evaluation_dataframe + +if __name__ == "__main__": + # STEP 1: Load MIMIC-III. dev=True keeps this to a small subset -- start + # here, and only drop it once the whole script runs end to end. + base_dataset = MIMIC3Dataset( + root="/srv/local/data/MIMIC-III/mimic-iii-clinical-database-1.4", + tables=["diagnoses_icd"], + dev=True, + ) + + # STEP 2: Apply the sequence-encoded generation task (no labels). + sample_dataset = base_dataset.set_task(EHRSequenceGenerationMIMIC3()) + print(f"Total samples: {len(sample_dataset)}") + + sample = sample_dataset[0] + # (num_visits, max_codes_per_visit) of vocabulary indices, right-padded + # with (0) -- NOT a multi-hot vector. + print(f"Visits tensor shape: {tuple(sample['visits'].shape)}") + + # STEP 3: Split by patient so no patient appears in two splits. + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + # STEP 4: Initialize GPT2 (small config for the dev subset). + model = GPT2( + dataset=sample_dataset, + embed_dim=128, + n_heads=4, + n_layers=4, + max_len=256, + batch_size=16, + epochs=5, + lr=1e-4, + save_dir="./gpt2_save", + ) + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params} parameters") + + # STEP 5: Train (saves the best checkpoint to save_dir). + model.train_model(train_dataset, val_dataset=val_dataset) + + # STEP 6: Generate one synthetic patient per real training patient. + synthetic = model.generate(num_samples=len(train_dataset)) + print("\nGenerated synthetic patients (first 3):") + for patient in synthetic[:3]: + print(f" {patient['patient_id']}: {len(patient['visits'])} visits") + print(f" {patient['visits']}") + + # STEP 7: Evaluate. The metrics want one row per (patient, visit, code); + # to_evaluation_dataframe produces exactly that from either real records or + # a generator's output. Real visits are read straight off the index tensor + # through the processor, which knows how to invert its own encoding. + processor = sample_dataset.input_processors["visits"] + index_to_code = {idx: code for code, idx in processor.code_vocab.items()} + + def to_records(subset): + for item in subset: + visits = [] + for row in item["visits"]: + codes = [ + index_to_code[idx] + for idx in processor.visit_code_ids(row) + if index_to_code.get(idx) not in (None, "", "") + ] + if codes: + visits.append(codes) + yield {"visits": visits} + + schema = {"visit_codes": str, "labels": int, "time": int, "id": str} + train_df = to_evaluation_dataframe(to_records(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(to_records(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) + print( + f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " + f"synthetic: {len(syn_df)}" + ) + + # Small settings for the dev subset; raise them on the full cohort. + results = evaluate_synthetic_ehr( + train_ehr=train_df, + test_ehr=test_df, + syn_ehr=syn_df, + sample_size=min(30, len(train_dataset), len(test_dataset)), + mode="lstm", + metrics="all", + lstm_params={"embed_dim": 16, "hidden_dim": 16, "batch_size": 16, "epochs": 3}, + n_bootstraps=5, + n_runs=3, + ) + print("\nGenerative metrics (mean +/- std):") + for name, (mean, std) in results.items(): + print(f" {name:30s} {mean:.4f} +/- {std:.4f}") diff --git a/examples/halo_mimic3.py b/examples/halo_mimic3.py index 3ea0a71be..1a2e11b80 100644 --- a/examples/halo_mimic3.py +++ b/examples/halo_mimic3.py @@ -2,19 +2,23 @@ This example demonstrates: 1. Loading MIMIC-III data -2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sequences) -3. Creating a SampleDataset with a NestedSequenceProcessor +2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sets, multi-hot + -- the encoding HALO consumes). See gpt2_mimic3.py for the sequence form + GPT2/PromptEHR need, and medgan_mimic3.py for the bag-of-codes form. +3. Creating a SampleDataset with a NestedMultiHotProcessor 4. Training the HALO generator with its custom training loop 5. Generating synthetic patients 6. Evaluating the synthetic data with the generative metrics suite """ -import pandas as pd - from pyhealth.datasets import MIMIC3Dataset, split_by_patient from pyhealth.metrics.generative import evaluate_synthetic_ehr from pyhealth.models import HALO -from pyhealth.tasks import EHRGenerationMIMIC3 +from pyhealth.tasks import ( + EHRGenerationMIMIC3, + decode_dataset, + to_evaluation_dataframe, +) if __name__ == "__main__": # STEP 1: Load MIMIC-III base dataset @@ -25,7 +29,8 @@ ) # STEP 2: Apply the EHR generation task (unconditional, no labels). - # This task is shared by all generators in pyhealth.models.generators. + # Extraction is shared across the generators; the encoding is not, so this + # is the multi-hot variant that HALO reads directly. sample_dataset = base_dataset.set_task(EHRGenerationMIMIC3()) print(f"Total samples: {len(sample_dataset)}") print(f"Input schema: {sample_dataset.input_schema}") @@ -34,6 +39,7 @@ sample = sample_dataset[0] print("\nSample structure:") print(f" Patient ID: {sample['patient_id']}") + # (num_visits, vocab_size) multi-hot -- there is no padded inner axis. print(f" Visits tensor shape: {tuple(sample['visits'].shape)}") # STEP 3: Split dataset by patient @@ -80,32 +86,16 @@ # train_df, test_df and syn_df below all share this exact schema. `labels` # is a placeholder here: privacy metrics ignore it and the utility metric # overwrites it with the next-visit prediction target. - index_to_code = { - v: k for k, v in sample_dataset.input_processors["visits"].code_vocab.items() - } - - def real_subset_to_records(subset): - for sample in subset: - pid = str(sample["patient_id"]) - visits_tensor = sample["visits"] - for t, visit in enumerate(visits_tensor.tolist()): - for idx in visit: - code = index_to_code.get(int(idx)) - if code in (None, "", ""): - continue - yield {"id": pid, "time": t, "visit_codes": code, "labels": 0} - - def synthetic_to_records(patients): - for p in patients: - pid = str(p["patient_id"]) - for t, visit in enumerate(p["visits"]): - for code in visit: - yield {"id": pid, "time": t, "visit_codes": code, "labels": 0} - + # Both conversions already live in the task module, so use them rather + # than re-deriving the encoding here: + # decode_dataset processed (or split) SampleDataset -> code records + # to_evaluation_dataframe records (real or synthetic) -> the long form + # Patients are renumbered 0, 1, 2, ... because synthetic patients do not + # correspond to real ones; the metrics only need a grouping key. schema = {"visit_codes": str, "labels": int, "time": int, "id": str} - train_df = pd.DataFrame(real_subset_to_records(train_dataset)).astype(schema) - test_df = pd.DataFrame(real_subset_to_records(test_dataset)).astype(schema) - syn_df = pd.DataFrame(synthetic_to_records(synthetic)).astype(schema) + train_df = to_evaluation_dataframe(decode_dataset(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(decode_dataset(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) print( f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " f"synthetic: {len(syn_df)}" diff --git a/examples/medgan_mimic3.py b/examples/medgan_mimic3.py new file mode 100644 index 000000000..4ff14a795 --- /dev/null +++ b/examples/medgan_mimic3.py @@ -0,0 +1,112 @@ +"""Example: train MedGAN on MIMIC-III and generate synthetic patients. + +MedGAN and CorGAN are bag-of-codes generators: a patient is one multi-hot +vector over the code vocabulary, with no visit axis at all. They read +EHRCodeSetGenerationMIMIC3, which pools every admission's codes into a single +set per patient -- not the per-visit tasks HALO (halo_mimic3.py) or +GPT2/PromptEHR (gpt2_mimic3.py) take. + +Swap MedGAN for CorGAN below and the rest of the script is unchanged. + +This example demonstrates: +1. Loading MIMIC-III data +2. Applying the EHRCodeSetGenerationMIMIC3 task (one ICD-9 code set per patient) +3. Training MedGAN (autoencoder pre-training, then adversarial training) +4. Generating synthetic patients +5. Evaluating with the privacy metrics + +Note on metrics: the utility metric in pyhealth.metrics.generative scores +next-visit prediction, which is meaningless without a visit axis. This example +therefore requests ``metrics="privacy"``. Asking for ``"all"`` here would +produce a utility number that looks real and is not. +""" + +from pyhealth.datasets import MIMIC3Dataset, split_by_patient +from pyhealth.metrics.generative import evaluate_synthetic_ehr +from pyhealth.models import MedGAN +from pyhealth.tasks import EHRCodeSetGenerationMIMIC3, to_evaluation_dataframe + +if __name__ == "__main__": + # STEP 1: Load MIMIC-III. dev=True keeps this small -- start here. + base_dataset = MIMIC3Dataset( + root="/srv/local/data/MIMIC-III/mimic-iii-clinical-database-1.4", + tables=["diagnoses_icd"], + dev=True, + ) + + # STEP 2: Pool each patient's admissions into one code set. min_visits + # still counts real admissions, so single-admission patients are dropped + # before the visit axis is collapsed. + sample_dataset = base_dataset.set_task(EHRCodeSetGenerationMIMIC3()) + print(f"Total samples: {len(sample_dataset)}") + + sample = sample_dataset[0] + # (vocab_size,) -- one row per patient, not per visit. + print(f"Visits tensor shape: {tuple(sample['visits'].shape)}") + + # STEP 3: Split by patient. + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + # STEP 4: Initialize MedGAN (small config for the dev subset). + model = MedGAN( + dataset=sample_dataset, + latent_dim=32, + hidden_dim=32, + discriminator_hidden_dim=64, + batch_size=32, + ae_epochs=10, + gan_epochs=20, + save_dir="./medgan_save", + ) + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params} parameters") + + # STEP 5: Train (autoencoder first, then the GAN). + model.train_model(train_dataset, val_dataset=val_dataset) + + # STEP 6: Generate one synthetic patient per real training patient. + synthetic = model.generate(num_samples=len(train_dataset)) + print("\nGenerated synthetic patients (first 3):") + for patient in synthetic[:3]: + # visits is a single-element list: the aggregate bag of codes. + print(f" {patient['patient_id']}: {len(patient['visits'][0])} codes") + + # STEP 7: Evaluate. MultiHotProcessor keeps its vocabulary in label_vocab + # (the nested processors call theirs code_vocab), and each row is a + # multi-hot vector, so the codes present are its nonzero columns. + processor = sample_dataset.input_processors["visits"] + index_to_code = {idx: code for code, idx in processor.label_vocab.items()} + + def to_records(subset): + for item in subset: + codes = [ + index_to_code[int(col)] + for col in item["visits"].nonzero(as_tuple=True)[0].tolist() + if index_to_code.get(int(col)) is not None + ] + # One "visit" per patient, matching the generator's output shape. + yield {"visits": [codes]} + + schema = {"visit_codes": str, "labels": int, "time": int, "id": str} + train_df = to_evaluation_dataframe(to_records(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(to_records(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) + print( + f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " + f"synthetic: {len(syn_df)}" + ) + + # privacy only -- see the note at the top of this file. + results = evaluate_synthetic_ehr( + train_ehr=train_df, + test_ehr=test_df, + syn_ehr=syn_df, + sample_size=min(30, len(train_dataset), len(test_dataset)), + metrics="privacy", + n_bootstraps=5, + ) + print("\nPrivacy metrics (mean +/- std):") + for name, (mean, std) in results.items(): + print(f" {name:30s} {mean:.4f} +/- {std:.4f}") diff --git a/pyhealth/models/generators/corgan.py b/pyhealth/models/generators/corgan.py index 807ed8066..0c1c64610 100644 --- a/pyhealth/models/generators/corgan.py +++ b/pyhealth/models/generators/corgan.py @@ -7,7 +7,8 @@ ``dataset -> SampleDataset -> model`` pipeline. CorGAN treats each patient as a flat bag-of-codes (no visit structure), so it -expects an input feature named ``visits`` backed by a ``MultiHotProcessor``. +expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, +which the :class:`~pyhealth.tasks.EHRCodeSetGenerationMIMIC3` task provides. Training has two phases (mirroring the reference): * a **convolutional autoencoder** is pre-trained with a sparse-friendly BCE @@ -350,7 +351,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "CorGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor." + "MultiHotProcessor (see EHRCodeSetGenerationMIMIC3)." ) self._batch_size = batch_size diff --git a/pyhealth/models/generators/gpt2.py b/pyhealth/models/generators/gpt2.py index 99b3bf92a..ac1d541ac 100644 --- a/pyhealth/models/generators/gpt2.py +++ b/pyhealth/models/generators/gpt2.py @@ -3,8 +3,10 @@ A simple decoder-only baseline that mirrors the standalone reference script ``generate_synthetic_mimic3_gpt2.py`` (``--mode transformer_baseline``) but plugged into the standard PyHealth ``dataset -> set_task -> SampleDataset -> -model`` pipeline. It consumes the same :class:`~pyhealth.tasks.EHRGeneration` -task as :class:`~pyhealth.models.HALO`. +model`` pipeline. It consumes the +:class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- the same extraction +:class:`~pyhealth.models.HALO` uses, but emitting code indices rather than +multi-hot rows, since a causal LM reads token ids. Each patient's visits are flattened into a single token stream:: @@ -15,7 +17,7 @@ (``do_sample`` + top-k/top-p) and decodes it back into per-visit code lists, splitting on the ``[VISIT_DELIM]`` token. -The code vocabulary is taken from the dataset's ``NestedSequenceProcessor`` +The code vocabulary is taken from the dataset's ``visits`` processor (which already reserves index 0 for ```` and index 1 for ````); three special tokens (BOS, EOS, VISIT_DELIM) are appended, and ```` (index 0) is reused as the padding token. @@ -39,8 +41,9 @@ class GPT2(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedSequenceProcessor}`` -- use the + :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- and whose + ``output_schema`` is empty. embed_dim: GPT-2 embedding dimension (``n_embd``). Must be divisible by ``n_heads``. Default: 512. n_heads: Number of attention heads. Default: 8. @@ -86,7 +89,17 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "GPT2 expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedSequenceProcessor (see EHRSequenceGenerationMIMIC3)." + ) + if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): + # Without this the visit row would be read as raw values. Under a + # multi-hot encoding every value is 1.0, so every code would silently + # become and training would look fine while learning nothing. + raise ValueError( + f"GPT2 needs a 'visits' processor that can invert its own " + f"encoding (a visit_code_ids method); got " + f"{type(dataset.input_processors['visits']).__name__}. Use " + "NestedSequenceProcessor, via the EHRSequenceGenerationMIMIC3 task." ) self.save_dir = save_dir @@ -95,7 +108,7 @@ def __init__( self._lr = lr self.max_len = max_len - # Code vocab from the NestedSequenceProcessor (includes =0, =1). + # Code vocab from the visits processor (includes =0, =1). self.visits_processor = dataset.input_processors["visits"] self.code_vocab_size = self.visits_processor.vocab_size() # Append three special tokens after the code vocab; reuse =0 as PAD. @@ -132,9 +145,9 @@ def _encode_visits(self, visits: torch.Tensor): """Flatten the padded visit-index tensor into causal-LM token streams. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. Index 0 is ```` and is - skipped. + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. + Index 0 is ```` and is skipped. Returns: input_ids: LongTensor ``(batch, L)`` token streams, right-padded. @@ -147,7 +160,7 @@ def _encode_visits(self, visits: torch.Tensor): n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) seq: List[int] = [self.bos_id] for j in range(n_visits): - codes = [int(c) for c in visits[i, j].tolist() if c > 0] + codes = self.visits_processor.visit_code_ids(visits[i, j]) seq.extend(codes) if j < n_visits - 1: seq.append(self.delim_id) @@ -176,8 +189,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/models/generators/halo.py b/pyhealth/models/generators/halo.py index 374c14000..2a51b59d2 100644 --- a/pyhealth/models/generators/halo.py +++ b/pyhealth/models/generators/halo.py @@ -365,14 +365,15 @@ class HALO(BaseModel): The model infers its code vocabulary from the fitted ``SampleDataset``: ``code_vocab_size = dataset.input_processors["visits"].vocab_size()`` - (the ``NestedSequenceProcessor`` vocab, which already reserves index 0 for + (the ``NestedMultiHotProcessor`` vocab, which already reserves index 0 for ```` and index 1 for ````). Three special tokens are appended for start-of-sequence, end-of-sequence, and pad-visit. Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedMultiHotProcessor}`` -- use the + :class:`~pyhealth.tasks.EHRGenerationMIMIC3` task -- and whose + ``output_schema`` is empty. embed_dim: Transformer embedding dimension (``n_embd``). Default: 768. n_heads: Number of attention heads. Must divide ``embed_dim``. Default: 12. @@ -420,7 +421,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "HALO expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedMultiHotProcessor (see EHRGenerationMIMIC3)." ) self.save_dir = save_dir @@ -428,7 +429,7 @@ def __init__( self._epochs = epochs self._lr = lr - # Code vocab from the NestedSequenceProcessor (includes , ). + # Code vocab from the NestedMultiHotProcessor (includes , ). self.visits_processor = dataset.input_processors["visits"] code_vocab_size = self.visits_processor.vocab_size() label_vocab_size = 0 # unconditional generation -- no output labels @@ -458,17 +459,25 @@ def __init__( # Multi-hot encoding helper # ------------------------------------------------------------------ def _encode_visits(self, visits: torch.Tensor): - """Convert a padded index tensor to HALO multi-hot format. + """Place per-visit multi-hot vectors into HALO's context window. - ``NestedSequenceProcessor`` returns code indices; the transformer - expects multi-hot vectors of shape ``(batch, n_ctx, total_vocab_size)`` - with special tokens. Layout (mirrors the reference): position 0 is the - start token, visits occupy positions 2+, the end token is placed on the - last visit's row, and the pad token fills the remaining positions. + Takes what :class:`~pyhealth.processors.NestedMultiHotProcessor` emits + -- one multi-hot row per visit -- and lays it out the way the + transformer expects: position 0 is the start token, visits occupy + positions 2+, the end token sits just past the last real visit, and the + pad token fills the rest. + + Fully vectorised, deliberately. This ran as a triple-nested Python loop + over (patient, visit, code slot) and cost ~108 ms per patient on an + A100 -- 99.8% of a training step, against ~0.03 s for the transformer's + own forward and backward. Two things made it that expensive: an + ``.item()`` per patient, each forcing a CUDA sync, and a separate + single-element kernel launch per code. Neither survives here. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)``. - Index 0 is ```` and is skipped. + visits: FloatTensor ``(batch, max_visits, code_vocab_size)``, + 1.0 where a code is present in that visit. A visit with no + codes is an all-zero row. Returns: batch_ehr: FloatTensor ``(batch, n_ctx, total_vocab_size)``. @@ -476,7 +485,8 @@ def _encode_visits(self, visits: torch.Tensor): with the autoregressive prediction targets. """ cfg = self.config - batch_size = visits.shape[0] + visits = visits.to(self.device) + batch_size, max_visits = visits.shape[0], visits.shape[1] batch_ehr = torch.zeros( batch_size, cfg.n_ctx, cfg.total_vocab_size, device=self.device @@ -487,19 +497,40 @@ def _encode_visits(self, visits: torch.Tensor): end_idx = start_idx + 1 pad_idx = start_idx + 2 - for i in range(batch_size): - # Count actual (non-padding) visits for this patient. - n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) - n_visits = min(n_visits, cfg.n_ctx - 2) - for j in range(n_visits): - for code_idx in visits[i, j]: - if code_idx > 0: # skip (index 0) - batch_ehr[i, j + 2, code_idx] = 1 - batch_mask[i, j + 2] = 1 + # Real visits per patient, for the whole batch at once. An all-zero row + # is an empty visit, exactly as a row of indices was before. + # + # TODO: this counts non-empty rows and then treats the *first* n_visits + # rows as the real ones, so a patient like [codes, empty, codes] would + # silently lose its last visit. The pre-vectorisation loop had the same + # behaviour, and EHRGenerationMIMIC3 cannot produce an interior + # empty visit (its __call__ skips codeless admissions), so nothing hits + # today -- but NestedMultiHotProcessor does emit all-zero rows for empty + # visits, so a hand-built SampleDataset can. Fix by masking on row + # position rather than row count. + n_visits = (visits.sum(dim=-1) > 0).sum(dim=1) + n_visits = n_visits.clamp(max=cfg.n_ctx - 2) # (batch,) + + # Two positions are reserved (start, end), so only this many visits fit. + keep = min(max_visits, cfg.n_ctx - 2) + if keep > 0: + pos = torch.arange(keep, device=self.device) # (keep,) + valid = (pos.unsqueeze(0) < n_visits.unsqueeze(1)) # (batch, keep) + # Codes occupy the first code_vocab_size columns; the three special + # tokens live above them and are set separately below. + batch_ehr[:, 2:2 + keep, :cfg.code_vocab_size] = ( + visits[:, :keep, :] * valid.unsqueeze(-1) + ) + batch_mask[:, 2:2 + keep, 0] = valid.to(batch_mask.dtype) + + batch_ehr[:, 0, start_idx] = 1 # start token + rows = torch.arange(batch_size, device=self.device) + batch_ehr[rows, n_visits + 1, end_idx] = 1 # end token - batch_ehr[i, 0, start_idx] = 1 # start token - batch_ehr[i, n_visits + 1, end_idx] = 1 # end token (on last visit) - batch_ehr[i, n_visits + 2:, pad_idx] = 1 # pad visits + # Everything past the end token is padding. + ctx = torch.arange(cfg.n_ctx, device=self.device) + is_pad = ctx.unsqueeze(0) >= (n_visits + 2).unsqueeze(1) # (batch, n_ctx) + batch_ehr[:, :, pad_idx] = is_pad.to(batch_ehr.dtype) batch_mask = batch_mask[:, 1:, :] # shift to align with shifted targets return batch_ehr, batch_mask @@ -511,8 +542,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: FloatTensor ``(batch, max_visits, code_vocab_size)`` from + the ``NestedMultiHotProcessor``, 1.0 where a code is present. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/models/generators/medgan.py b/pyhealth/models/generators/medgan.py index 3e0ee06c0..6faea9214 100644 --- a/pyhealth/models/generators/medgan.py +++ b/pyhealth/models/generators/medgan.py @@ -7,7 +7,8 @@ ``dataset -> SampleDataset -> model`` pipeline. MedGAN treats each patient as a flat bag-of-codes (no visit structure), so it -expects an input feature named ``visits`` backed by a ``MultiHotProcessor``. +expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, +which the :class:`~pyhealth.tasks.EHRCodeSetGenerationMIMIC3` task provides. The training procedure has two phases (mirroring the reference): * a **linear autoencoder** is pre-trained with binary cross-entropy @@ -231,7 +232,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "MedGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor." + "MultiHotProcessor (see EHRCodeSetGenerationMIMIC3)." ) # The generator's residual connection (``out + residual`` with diff --git a/pyhealth/models/generators/promptehr.py b/pyhealth/models/generators/promptehr.py index 4622e72c6..6c50690d4 100644 --- a/pyhealth/models/generators/promptehr.py +++ b/pyhealth/models/generators/promptehr.py @@ -3,7 +3,7 @@ This is a PyHealth ``BaseModel`` port of PromptEHR (Wang & Sun, EMNLP'22, https://github.com/RyanWangZf/PromptEHR), wrapped so it consumes the standard ``dataset -> set_task -> SampleDataset -> model`` pipeline and shares the same -:class:`~pyhealth.tasks.EHRGeneration` task as +:class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task as :class:`~pyhealth.models.HALO` and :class:`~pyhealth.models.GPT2`. PromptEHR treats sequential EHRs as a *neural database* and learns to fill in @@ -17,7 +17,7 @@ the way :class:`~pyhealth.models.GPT2` wraps ``GPT2LMHeadModel``. * **Prompt learning.** The reference reparameterizes a learnable prompt from patient baseline demographics and prepends it to the encoder/decoder - (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.EHRGeneration` + (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task is *unconditional* (only ``visits``, no baseline features -- exactly like HALO/GPT2), so the prompt reduces to a learnable continuous **soft prefix** prepended to the encoder. This is the prompt-tuning core without the @@ -35,10 +35,10 @@ [CODE_PROMPT] [VISIT_DELIM] ... [EOS] The reference handles several code types (diagnosis / procedure / drug / lab) -each with its own modality prompt token; the PyHealth ``EHRGeneration`` task +each with its own modality prompt token; the PyHealth ``EHRSequenceGenerationMIMIC3`` task exposes a single ``visits`` modality, so a single ``[CODE_PROMPT]`` token marks it. The code vocabulary is taken from the dataset's -``NestedSequenceProcessor`` (which already reserves index 0 for ```` and +``visits`` processor (which already reserves index 0 for ```` and index 1 for ````); five special tokens (BOS, EOS, VISIT_DELIM, MASK, CODE_PROMPT) are appended, and ```` (index 0) is reused as the pad token. """ @@ -63,12 +63,13 @@ class PromptEHR(BaseModel): Trains a BART denoising autoencoder with a learnable soft prompt on patient visit-code streams, then generates synthetic patients by prompt-conditioned encoder-decoder sampling. Generation is **unconditional** (no demographic - conditioning), matching the :class:`~pyhealth.tasks.EHRGeneration` task. + conditioning), matching the :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task. Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedSequenceProcessor}`` -- use the + :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- and whose + ``output_schema`` is empty. embed_dim: BART model dimension (``d_model``). Must be divisible by ``n_heads``. Default: 256. n_heads: Number of attention heads (encoder and decoder). Default: 8. @@ -127,7 +128,17 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "PromptEHR expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedSequenceProcessor (see EHRSequenceGenerationMIMIC3)." + ) + if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): + # Without this the visit row would be read as raw values. Under a + # multi-hot encoding every value is 1.0, so every code would silently + # become and training would look fine while learning nothing. + raise ValueError( + f"PromptEHR needs a 'visits' processor that can invert its own " + f"encoding (a visit_code_ids method); got " + f"{type(dataset.input_processors['visits']).__name__}. Use " + "NestedSequenceProcessor, via the EHRSequenceGenerationMIMIC3 task." ) self.save_dir = save_dir @@ -139,7 +150,7 @@ def __init__( self.mean_span_len = mean_span_len self.prompt_length = prompt_length - # Code vocab from the NestedSequenceProcessor (includes =0, =1). + # Code vocab from the visits processor (includes =0, =1). self.visits_processor = dataset.input_processors["visits"] self.code_vocab_size = self.visits_processor.vocab_size() # Append five special tokens after the code vocab; reuse =0 as PAD. @@ -197,7 +208,7 @@ def _serialize(self, visits: torch.Tensor) -> List[List[int]]: n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) seq: List[int] = [self.code_prompt_id] for j in range(n_visits): - codes = [int(c) for c in visits[i, j].tolist() if c > 0] + codes = self.visits_processor.visit_code_ids(visits[i, j]) seq.extend(codes) if j < n_visits - 1: seq.append(self.delim_id) @@ -313,8 +324,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass (denoising seq2seq reconstruction). Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..5fd0a12ba 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -26,6 +26,7 @@ def get_processor(name: str): RegressionLabelProcessor, ) from .multi_hot_processor import MultiHotProcessor +from .nested_multihot_processor import NestedMultiHotProcessor from .nested_sequence_processor import ( NestedFloatsProcessor, NestedSequenceProcessor, @@ -66,6 +67,7 @@ def get_processor(name: str): "LabelProcessor", "MultiHotProcessor", "NestedFloatsProcessor", + "NestedMultiHotProcessor", "NestedSequenceProcessor", "RawProcessor", "SequenceProcessor", diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py new file mode 100644 index 000000000..3b6c7e10e --- /dev/null +++ b/pyhealth/processors/nested_multihot_processor.py @@ -0,0 +1,176 @@ +from collections.abc import Iterable +from typing import Any + +import torch + +from . import register_processor +from .base_processor import FeatureProcessor, TokenProcessorInterface + + +@register_processor("nested_multihot") +class NestedMultiHotProcessor(FeatureProcessor, TokenProcessorInterface): + """Nested categorical sequences as per-visit multi-hot vectors. + + Same input as :class:`NestedSequenceProcessor` -- a list of visits, each a + list of codes -- but it emits what set-membership models actually consume: + one row per visit, one column per vocabulary entry, 1 where the code is + present. + + Why this exists rather than indices + ----------------------------------- + ``NestedSequenceProcessor`` pads every visit to the longest visit seen + during ``fit``. That width is set by a single outlier: on eICU one visit + holds 3,951 code entries (the same diagnosis re-charted through a stay) + while a typical visit holds about five. Every downstream consumer then pays + for it twice -- once in memory, once again if it loops over the padding to + find the handful of real codes. + + For a 921-code vocabulary on eICU the multi-hot form is **8.6x smaller** + than the padded index form (4.5 KB vs 38.7 KB per patient), because it is + sized by the vocabulary rather than by the worst-case visit. + + Repeats collapse + ---------------- + A code charted five times in one visit sets the same column to 1 once. That + matches what set-membership models already do with the index form, so + swapping this in changes nothing about what a model sees. If you need + counts, this is the wrong processor. + + Special tokens + -------------- + ```` (0) and ```` (1) keep their indices, so the vocabulary is + interchangeable with ``NestedSequenceProcessor`` -- a cached vocabulary from + one restores into the other. Column 0 is therefore always zero: ```` + means "nothing here", and marking it would make padding indistinguishable + from a real code. + + Examples: + >>> processor = NestedMultiHotProcessor() + >>> samples = [{"codes": [["A", "B"], ["C"]]}] + >>> processor.fit(samples, "codes") + >>> out = processor.process([["A", "B"], ["A"]]) + >>> out.shape # (2 visits, vocab_size) + torch.Size([2, 5]) + >>> out[0].nonzero().flatten().tolist() # A and B present in visit 0 + [2, 3] + """ + + def __init__(self, padding: int = 0): + # `padding` is accepted and ignored so this is a drop-in swap for + # NestedSequenceProcessor in a schema. There is no inner axis to pad -- + # that is the entire point -- so honouring it would be misleading. + self.code_vocab: dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 + self._padding = padding + + def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: + """Build the vocabulary. Inner length is irrelevant here, so unlike + ``NestedSequenceProcessor`` nothing is measured about visit width. + + Args: + samples: Sample dictionaries. + field: Field holding the nested sequence. + """ + for sample in samples: + if field not in sample or sample[field] is None: + continue + nested_seq = sample[field] + if not isinstance(nested_seq, list): + continue + for inner_seq in nested_seq: + if not isinstance(inner_seq, list): + continue + for code in inner_seq: + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 + + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + self._next_index = len(self.code_vocab) + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + + def process(self, value: list[list[Any]]) -> torch.Tensor: + """Nested sequence -> ``(num_visits, vocab_size)`` float multi-hot. + + Built with one ``scatter_`` per visit rather than per-code assignment, + so cost tracks the number of real codes, never the vocabulary size or a + padded width. + + An empty or ``None`` sample yields a single all-zero visit, mirroring + ``NestedSequenceProcessor`` returning one all-```` row: both say + "one visit, nothing in it". + + Args: + value: Nested list of codes ``[[code1, code2], [code3], ...]``. + + Returns: + 2D float tensor ``(num_visits, vocab_size)``, 1.0 where present. + """ + vocab_size = len(self.code_vocab) + unk = self.code_vocab[""] + + if not value or len(value) == 0: + return torch.zeros(1, vocab_size, dtype=torch.float) + + out = torch.zeros(len(value), vocab_size, dtype=torch.float) + for row, inner_seq in enumerate(value): + if inner_seq is None or len(inner_seq) == 0: + continue # empty visit stays all-zero + idx = [self.code_vocab.get(code, unk) if code is not None else unk + for code in inner_seq] + # scatter_ over the visit's own codes: duplicates write 1.0 twice, + # which is still 1.0 -- repeats collapse by construction. + out[row].scatter_(0, torch.tensor(idx, dtype=torch.long), + torch.ones(len(idx))) + return out + + def size(self) -> int: + """Feature width: the vocabulary, since that is the row length.""" + return len(self.code_vocab) + + def vocab_size(self) -> int: + """Return vocabulary size.""" + return len(self.code_vocab) + + def __repr__(self): + return f"NestedMultiHotProcessor(vocab_size={len(self.code_vocab)})" + + def is_token(self) -> bool: + """Output is a dense indicator vector, not token indices.""" + return False + + def schema(self) -> tuple[str, ...]: + return ("value",) + + def dim(self) -> tuple[int, ...]: + """Output is a 2D tensor (visits, vocab).""" + return (2,) + + def spatial(self) -> tuple[bool, ...]: + # Visits (time) are ordered; the vocabulary axis is an unordered set. + return (True, False) diff --git a/pyhealth/processors/nested_sequence_processor.py b/pyhealth/processors/nested_sequence_processor.py index c03c800d0..060b5e3f8 100644 --- a/pyhealth/processors/nested_sequence_processor.py +++ b/pyhealth/processors/nested_sequence_processor.py @@ -154,6 +154,29 @@ def vocab_size(self) -> int: """Return the size of the processor's vocabulary.""" return len(self.code_vocab) + def visit_code_ids(self, row: torch.Tensor) -> list[int]: + """Code indices present in one processed visit row. + + The inverse of what :meth:`process` writes, for consumers that need a + code *list* rather than the tensor -- the sequence generators GPT2 and + PromptEHR, and :func:`pyhealth.tasks.decode_dataset`. Each nested + processor implements this for its own encoding, so a consumer never has + to guess how to read a row. + + Here the row already holds indices, right-padded with ```` (0). + + Args: + row: 1D tensor, one processed visit. + + Returns: + Code indices in charted order, ```` dropped. + + Examples: + >>> processor.visit_code_ids(torch.tensor([2, 4, 0])) + [2, 4] + """ + return [int(c) for c in row.tolist() if c > 0] + def size(self) -> int: """Return max inner length (embedding dimension) for unified API.""" return self._max_inner_len diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 612bc2546..2e7bd6bab 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -52,9 +52,12 @@ MortalityPredictionStageNetMIMIC4, ) from .generate_ehr import ( - EHRGeneration, + EHRCodeSetGenerationMIMIC3 as EHRCodeSetGenerationMIMIC3, + EHRCodeSetGenerationMIMIC4 as EHRCodeSetGenerationMIMIC4, EHRGenerationMIMIC3, EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3 as EHRSequenceGenerationMIMIC3, + EHRSequenceGenerationMIMIC4 as EHRSequenceGenerationMIMIC4, decode_dataset, to_evaluation_dataframe, ) diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 6fb23da9d..3f7fcd99d 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -1,14 +1,30 @@ """EHR sequence-generation tasks for PyHealth generative models. -This is the shared task for every generator in -:mod:`pyhealth.models.generators` (HALO, MedGAN, CorGAN, PromptEHR, ...). It -extracts, for each patient, the ordered list of visits where each visit is the -list of medical codes recorded in that admission. The single input feature -``visits`` is processed by :class:`~pyhealth.processors.NestedSequenceProcessor`; -there is no prediction label, so ``output_schema`` is empty. - -:class:`EHRGeneration` holds all the extraction logic; dataset-specific -subclasses only declare which event type and code attribute to read. +These back every generator in :mod:`pyhealth.models.generators`. Each one +extracts, for each patient, the codes recorded across their admissions. There +is no prediction label, so ``output_schema`` is empty. + +The classes are flat and independent -- one per (model family, dataset) -- with +the dataset named in the class, because the extraction is MIMIC-shaped: it +assumes an ``admissions`` event type and a ``hadm_id`` linking codes to an +admission. eICU, OMOP and MEDS do not look like that, so a task for those +datasets belongs alongside these rather than inheriting from them. + +============================= ========================== ================== +Task Encoding Models +============================= ========================== ================== +``EHRGenerationMIMIC3/4`` per-visit multi-hot rows HALO +``EHRSequenceGenerationMIMIC3/4`` per-visit code indices GPT2, PromptEHR +``EHRCodeSetGenerationMIMIC3/4`` one code set per patient MedGAN, CorGAN +============================= ========================== ================== + +Match the task to the model. Handing a model the wrong encoding does not raise: +the numbers still have the right shape and dtype, so training runs and produces +a confidently wrong result. That is why these are separate classes rather than +one task with a flag. + +``event_type`` / ``code_attr`` are class attributes naming the dataset's coding +columns; override them on an instance to read a different table. Evaluating generated data ------------------------- @@ -63,103 +79,268 @@ """ import logging -from typing import Callable, Dict, List, Optional, Type, Union +from collections.abc import Callable +from typing import ClassVar from pyhealth.data.data import Patient -from pyhealth.processors import NestedSequenceProcessor +from pyhealth.processors import ( + MultiHotProcessor, + NestedMultiHotProcessor, + NestedSequenceProcessor, +) from .base_task import BaseTask logger = logging.getLogger(__name__) -class EHRGeneration(BaseTask): - """Generic per-visit code-sequence task for unconditional EHR generators. - - Builds one sample per qualifying patient: the ordered list of visits, each - visit being the list of codes (read from ``code_attr`` on ``event_type`` - events) recorded in that admission. Patients with fewer than ``min_visits`` - qualifying visits are skipped. +def _mimic_visits( + patient: Patient, event_type: str, code_attr: str +) -> list[list[str]]: + """Ordered per-admission code lists for a MIMIC-style patient. - Subclass and override the class attributes for a specific dataset, or set - them on an instance. The defaults read MIMIC-III ICD-9 diagnosis codes. + Deliberately MIMIC-specific and deliberately not a base-class method: it + assumes an ``admissions`` event type and a ``hadm_id`` linking codes to an + admission. eICU, OMOP and MEDS do not look like this, so a task for those + writes its own extraction rather than inheriting one that would silently + return nothing. Args: - task_name: Name of the task. - input_schema: ``{"visits": NestedSequenceProcessor}``. - output_schema: empty (generative task, no labels). - event_type: Event type to pull per admission. Default - ``"diagnoses_icd"``. - code_attr: Event attribute holding the code string. Default - ``"icd9_code"``. - min_visits: Minimum qualifying visits to keep a patient. Default 2. + patient: Patient to read. + event_type: Event type carrying the codes (e.g. ``"diagnoses_icd"``). + code_attr: Attribute on those events holding the code string. + + Returns: + One list of codes per admission, admissions with no codes dropped. """ + visits: list[list[str]] = [] + for admission in patient.get_events(event_type="admissions"): + events = patient.get_events( + event_type=event_type, + filters=[("hadm_id", "==", admission.hadm_id)], + ) + codes = [ + getattr(event, code_attr) + for event in events + if getattr(event, code_attr, None) + ] + if codes: + visits.append(codes) + return visits + + +class EHRGenerationMIMIC3(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. - task_name: str = "ehr_generation" - input_schema: Dict[str, Union[str, Type]] = {"visits": NestedSequenceProcessor} - output_schema: Dict[str, Union[str, Type]] = {} + 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": NestedMultiHotProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" min_visits: int = 2 - def __call__(self, patient: Patient) -> List[Dict]: + def __call__(self, patient: Patient) -> list[dict]: """Extract the per-visit code sequence for a patient.""" - visits: List[List[str]] = [] - admissions = patient.get_events(event_type="admissions") - for admission in admissions: - events = patient.get_events( - event_type=self.event_type, - filters=[("hadm_id", "==", admission.hadm_id)], - ) - codes = [ - getattr(event, self.code_attr) - for event in events - if getattr(event, self.code_attr, None) - ] - if codes: - visits.append(codes) + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + return [{"patient_id": patient.patient_id, "visits": visits}] + + +class EHRGenerationMIMIC4(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 :class:`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": NestedMultiHotProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} + + event_type: str = "diagnoses_icd" + code_attr: str = "icd_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + return [{"patient_id": patient.patient_id, "visits": visits}] + + +class EHRSequenceGenerationMIMIC3(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": NestedSequenceProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} + + event_type: str = "diagnoses_icd" + code_attr: str = "icd9_code" + min_visits: int = 2 + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) if len(visits) < self.min_visits: return [] + return [{"patient_id": patient.patient_id, "visits": visits}] + + +class EHRSequenceGenerationMIMIC4(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": NestedSequenceProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" + code_attr: str = "icd_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] return [{"patient_id": patient.patient_id, "visits": visits}] -class EHRGenerationMIMIC3(EHRGeneration): - """EHR generation task for MIMIC-III (ICD-9 diagnosis codes). +class EHRCodeSetGenerationMIMIC3(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 + :mod:`pyhealth.metrics.generative` is not meaningful for these models; + see this module's header. Examples: >>> from pyhealth.datasets import MIMIC3Dataset - >>> from pyhealth.tasks import EHRGenerationMIMIC3 + >>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC3 >>> dataset = MIMIC3Dataset( - ... root="/path/to/mimic-iii/1.4", - ... tables=["diagnoses_icd"], + ... root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"] ... ) - >>> samples = dataset.set_task(EHRGenerationMIMIC3()) + >>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC3()) + >>> samples[0]["visits"].shape # (vocab_size,) + torch.Size([512]) """ - task_name: str = "ehr_generation_mimic3" + task_name: str = "ehr_codeset_generation_mimic3" + input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Pool every visit's codes into one per-patient set.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + codes = sorted({code for visit in visits for code in visit}) + return [{"patient_id": patient.patient_id, "visits": codes}] -class EHRGenerationMIMIC4(EHRGeneration): - """EHR generation task for MIMIC-IV (ICD diagnosis codes). +class EHRCodeSetGenerationMIMIC4(BaseTask): + """One pooled ICD code set per MIMIC-IV patient. For MedGAN and CorGAN. Examples: >>> from pyhealth.datasets import MIMIC4Dataset - >>> from pyhealth.tasks import EHRGenerationMIMIC4 + >>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC4 >>> dataset = MIMIC4Dataset( ... ehr_root="/path/to/mimiciv/2.2/", ... ehr_tables=["patients", "admissions", "diagnoses_icd"], ... ) - >>> samples = dataset.set_task(EHRGenerationMIMIC4()) + >>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC4()) + >>> samples[0]["visits"].shape # (vocab_size,) + torch.Size([512]) """ - task_name: str = "ehr_generation_mimic4" + task_name: str = "ehr_codeset_generation_mimic4" + input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" code_attr: str = "icd_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Pool every visit's codes into one per-patient set.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + codes = sorted({code for visit in visits for code in visit}) + return [{"patient_id": patient.patient_id, "visits": codes}] # ---------------------------------------------------------------------------- @@ -167,7 +348,7 @@ class EHRGenerationMIMIC4(EHRGeneration): # ---------------------------------------------------------------------------- def to_evaluation_dataframe( records, - label_fn: Optional[Callable[[Dict], int]] = None, + label_fn: Callable[[dict], int] | None = None, subject_col: str = "id", visit_col: str = "time", code_col: str = "visit_codes", @@ -185,7 +366,7 @@ def to_evaluation_dataframe( Args: records: Iterable of ``{"visits": [[code, ...], ...]}`` dicts. Both the - :class:`EHRGeneration` task output and a generator's ``generate()`` + generation tasks' output and a generator's ``generate()`` output have this shape. label_fn: Optional callable mapping a record to a binary patient label (0/1) used by the utility metrics. Defaults to all-zeros. @@ -197,6 +378,15 @@ def to_evaluation_dataframe( 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 """ import pandas as pd @@ -218,34 +408,60 @@ def to_evaluation_dataframe( ) -def decode_dataset(sample_dataset, feature_key: str = "visits") -> List[Dict]: - """Decode a processed EHRGeneration ``SampleDataset`` back into records. +def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: + """Decode a processed multi-hot ``SampleDataset`` back into code records. - Inverts the :class:`~pyhealth.processors.NestedSequenceProcessor` encoding + Inverts the :class:`~pyhealth.processors.NestedMultiHotProcessor` encoding using its vocabulary (skipping ````/````), 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. + Args: - sample_dataset: A ``SampleDataset`` produced by :class:`EHRGeneration`. + sample_dataset: A ``SampleDataset`` (or a split of one) produced by + :class:`EHRGenerationMIMIC3` / :class:`EHRGenerationMIMIC4`. feature_key: Input feature key holding the nested code sequence. Default ``"visits"``. Returns: List of ``{"visits": [[code_str, ...], ...]}`` records. + + Raises: + TypeError: If ``feature_key`` is not backed by a + :class:`~pyhealth.processors.NestedMultiHotProcessor`. + + Examples: + >>> from pyhealth.tasks.generate_ehr import decode_dataset + >>> records = decode_dataset(samples) + >>> records[0]["visits"][0] + ['4019', '25000'] """ - processor = sample_dataset.input_processors[feature_key] + # split_by_patient hands back a torch Subset, which carries no processors + # of its own -- decoding a split is the common case, so resolve through it. + source = getattr(sample_dataset, "dataset", sample_dataset) + processor = source.input_processors[feature_key] + if not isinstance(processor, NestedMultiHotProcessor): + raise TypeError( + f"decode_dataset inverts the multi-hot encoding, but '{feature_key}' " + f"is a {type(processor).__name__}. Use one of the multi-hot tasks\n (EHRGenerationMIMIC3/4), or " + "read the codes off the index tensor directly." + ) index_to_code = {idx: code for code, idx in processor.code_vocab.items()} - records: List[Dict] = [] + records: list[dict] = [] for i in range(len(sample_dataset)): sample = sample_dataset[i] - visits: List[List[str]] = [] - for row in sample[feature_key].tolist(): + visits: list[list[str]] = [] + # Each row is a multi-hot vector over the vocabulary, so the codes + # present are its nonzero columns -- the values are all 1.0. + for row in sample[feature_key]: codes = [ - index_to_code[int(idx)] - for idx in row - if index_to_code.get(int(idx)) not in (None, "", "") + index_to_code[int(col)] + for col in row.nonzero(as_tuple=True)[0].tolist() + if index_to_code.get(int(col)) not in (None, "", "") ] if codes: visits.append(codes) diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py new file mode 100644 index 000000000..fee993523 --- /dev/null +++ b/tests/core/test_generator_encodings.py @@ -0,0 +1,228 @@ +"""Each generator family gets the encoding it actually consumes. + +Extraction is shared across the EHR-generation tasks, the encoding is not: +HALO reads multi-hot rows, GPT2/PromptEHR read code indices, MedGAN/CorGAN read +one pooled set per patient. Pairing a model with the wrong task does not raise +-- the shapes and dtypes stay valid -- so these tests pin the pairing itself, +and check that real codes survive the round trip rather than collapsing to +. +""" + +import unittest +from typing import ClassVar + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import GPT2, PromptEHR +from pyhealth.processors import ( + MultiHotProcessor, + NestedMultiHotProcessor, + NestedSequenceProcessor, +) +from pyhealth.tasks import ( + EHRCodeSetGenerationMIMIC3, + EHRCodeSetGenerationMIMIC4, + EHRGenerationMIMIC3, + EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, + EHRSequenceGenerationMIMIC4, +) +from pyhealth.tasks.base_task import BaseTask + +SAMPLES = [ + {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, + {"patient_id": "p1", "visits": [["A05B"], ["A04A", "B035"]]}, + {"patient_id": "p2", "visits": [["C129", "A11D"], ["A05C"], ["A04A"]]}, + {"patient_id": "p3", "visits": [["B035"], ["A05B", "C129"]]}, +] + + +def _dataset(schema_key, name): + return create_sample_dataset( + samples=SAMPLES, + input_schema={"visits": schema_key}, + output_schema={}, + dataset_name=name, + ) + + +class _Event: + def __init__(self, hadm_id, code=None): + self.hadm_id = hadm_id + self.icd9_code = code + + +class _Patient: + """Minimal stand-in for pyhealth.data.Patient: admissions plus coded events.""" + + def __init__(self, patient_id, visits): + self.patient_id = patient_id + self._admissions = [_Event(f"h{i}") for i in range(len(visits))] + self._codes = [ + _Event(f"h{i}", code) + for i, codes in enumerate(visits) + for code in codes + ] + + def get_events(self, event_type, filters=None): + if event_type == "admissions": + return self._admissions + hadm = filters[0][2] + return [e for e in self._codes if e.hadm_id == hadm] + + +class TestExtraction(unittest.TestCase): + """Extraction, and the pooling the code-set tasks add.""" + + VISITS: ClassVar[list] = [["A05B", "A05C"], ["A11D"], ["A05B"]] + + def test_per_visit_tasks_keep_visit_structure(self): + patient = _Patient("p0", self.VISITS) + samples = EHRGenerationMIMIC3()(patient) + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["visits"], self.VISITS) + # Same extraction regardless of encoding -- only input_schema differs. + self.assertEqual(EHRSequenceGenerationMIMIC3()(patient)[0]["visits"], + self.VISITS) + + def test_codeset_task_pools_and_dedupes(self): + samples = EHRCodeSetGenerationMIMIC3()(_Patient("p0", self.VISITS)) + self.assertEqual(len(samples), 1) + # One flat set: A05B appears in two visits and survives once. + self.assertEqual(samples[0]["visits"], ["A05B", "A05C", "A11D"]) + + def test_min_visits_counts_real_visits_before_pooling(self): + """Pooling must not let a 1-visit patient past a min_visits=2 filter.""" + one_visit = _Patient("p1", [["A05B", "A05C", "A11D"]]) + self.assertEqual(EHRCodeSetGenerationMIMIC3()(one_visit), []) + self.assertEqual(EHRGenerationMIMIC3()(one_visit), []) + + def test_codeless_admissions_are_dropped(self): + patient = _Patient("p2", [["A05B"], [], ["A11D"]]) + self.assertEqual(EHRGenerationMIMIC3()(patient)[0]["visits"], + [["A05B"], ["A11D"]]) + + +class TestTaskEncodings(unittest.TestCase): + """Each task declares the processor its models consume.""" + + def test_each_family_gets_its_own_encoding(self): + for task in (EHRGenerationMIMIC3, EHRGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], NestedMultiHotProcessor) + for task in (EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], NestedSequenceProcessor) + for task in (EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], MultiHotProcessor) + + def test_tasks_are_flat(self): + """No task inherits from another: the MIMIC extraction is not a base. + + A parent class would invite subclassing it for eICU/OMOP, where there + is no ``admissions`` event type and no ``hadm_id`` -- which would + return no samples rather than fail. + """ + tasks = [ + EHRGenerationMIMIC3, EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4, + EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4, + ] + for task in tasks: + with self.subTest(task=task.__name__): + self.assertEqual(task.__bases__, (BaseTask,)) + self.assertIn("MIMIC", task.__name__) + + def test_task_names_are_unique(self): + tasks = [ + EHRGenerationMIMIC3, EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4, + EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4, + ] + names = [t.task_name for t in tasks] + self.assertEqual(len(names), len(set(names))) + + def test_mimic4_reads_the_mimic4_code_column(self): + self.assertEqual(EHRGenerationMIMIC3.code_attr, "icd9_code") + self.assertEqual(EHRGenerationMIMIC4.code_attr, "icd_code") + + +class TestVisitCodeIds(unittest.TestCase): + """NestedSequenceProcessor inverts its own rows for the token generators.""" + + def test_matches_the_multihot_columns(self): + """Both encodings of the same visit name the same codes. + + NestedMultiHotProcessor has no visit_code_ids -- nothing consumes one -- + so its codes are read here the way decode_dataset reads them, as the + row's nonzero columns. + """ + multihot = _dataset("nested_multihot", "vci_mh") + indexed = _dataset("nested_sequence", "vci_ix") + ix_proc = indexed.input_processors["visits"] + + # Same samples, same traversal order, so the vocabularies must match -- + # that is what makes the per-visit comparison below meaningful. + self.assertEqual(multihot.input_processors["visits"].code_vocab, + ix_proc.code_vocab) + + for i in range(len(SAMPLES)): + mh_row = multihot[i]["visits"] + ix_row = indexed[i]["visits"] + for visit in range(mh_row.shape[0]): + # Multi-hot columns come out in vocabulary order, the index form + # in charted order, so compare as sets. + self.assertEqual( + set(mh_row[visit].nonzero(as_tuple=True)[0].tolist()), + set(ix_proc.visit_code_ids(ix_row[visit])), + ) + + def test_padding_is_dropped_not_read_as_a_code(self): + processor = _dataset("nested_sequence", "vci_pad").input_processors["visits"] + row = processor.process([["A05B", "A05C"]])[0] + ids = processor.visit_code_ids(row) + self.assertEqual(len(ids), 2) + self.assertNotIn(processor.PAD, ids) + self.assertNotIn(processor.UNK, ids) + + +class TestTokenGeneratorsOnIndices(unittest.TestCase): + """GPT2 and PromptEHR serialise real codes from their own encoding.""" + + MODELS: ClassVar[list] = [ + (GPT2, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64}), + (PromptEHR, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64, + "prompt_length": 4}), + ] + + def _streams(self, model, visits): + if hasattr(model, "_serialize"): + return model._serialize(visits) + input_ids, _, _ = model._encode_visits(visits) + return [row.tolist() for row in input_ids] + + def test_codes_survive_serialisation(self): + for cls, kwargs in self.MODELS: + with self.subTest(model=cls.__name__): + dataset = _dataset("nested_sequence", f"gen_{cls.__name__}") + model = cls(dataset=dataset, batch_size=2, epochs=1, **kwargs) + # Patients have different visit counts, so let the dataloader + # pad the visit dimension rather than stacking raw samples. + batch = next(iter(get_dataloader(dataset, batch_size=2))) + visits = batch["visits"] + streams = self._streams(model, visits) + + code_ids = [ + t + for stream in streams + for t in stream + if t < model.code_vocab_size and t != 0 + ] + self.assertTrue(code_ids, "no code tokens were emitted at all") + self.assertGreater( + len(set(code_ids) - {model.visits_processor.UNK}), + 1, + f"{cls.__name__} emitted only : the visit row was read " + "as raw values instead of via visit_code_ids", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_halo.py b/tests/core/test_halo.py index be9788d62..f4abe96fc 100644 --- a/tests/core/test_halo.py +++ b/tests/core/test_halo.py @@ -20,7 +20,7 @@ def setUp(self): ] # Generative task: one nested-sequence input feature, no output labels. - self.input_schema = {"visits": "nested_sequence"} + self.input_schema = {"visits": "nested_multihot"} self.output_schema = {} self.dataset = create_sample_dataset( @@ -57,7 +57,30 @@ def test_forward_input_format(self): loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) self.assertIsInstance(batch["visits"], torch.Tensor) - self.assertEqual(batch["visits"].dim(), 3) # (B, max_visits, max_codes) + self.assertEqual(batch["visits"].dim(), 3) # (B, max_visits, vocab_size) + + def test_encode_truncates_past_context(self): + """A patient with more visits than n_ctx-2 is cut, not wrapped.""" + many = [{"patient_id": "long", "visits": [["A05B"]] * 20}] + dataset = create_sample_dataset( + samples=many + self.samples, + input_schema=self.input_schema, + output_schema=self.output_schema, + dataset_name="test_halo_long", + ) + model = HALO( + dataset=dataset, + embed_dim=16, + n_heads=2, + n_layers=2, + n_ctx=8, + batch_size=2, + epochs=1, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + ehr, mask = model._encode_visits(batch["visits"]) + self.assertEqual(ehr.shape[1], 8) + self.assertEqual(mask.shape[1], 7) def test_model_forward(self): """Forward returns a finite scalar loss and a probability tensor.""" diff --git a/tests/core/test_nested_multihot_processor.py b/tests/core/test_nested_multihot_processor.py new file mode 100644 index 000000000..a83c485f9 --- /dev/null +++ b/tests/core/test_nested_multihot_processor.py @@ -0,0 +1,94 @@ +"""Tests for NestedMultiHotProcessor. + +The processor exists to feed set-membership models (HALO) without materialising +the padded index form, whose width is set by the single longest visit in the +dataset. These tests pin the properties that make it a safe swap: the vocabulary +lays out identically to NestedSequenceProcessor, repeats collapse, and empty +visits stay empty. +""" + +import unittest + +import torch + +from pyhealth.processors import NestedMultiHotProcessor, NestedSequenceProcessor + + +class TestNestedMultiHotProcessor(unittest.TestCase): + def setUp(self): + self.samples = [ + {"codes": [["A", "B"], ["C"]]}, + {"codes": [["B", "D", "E"]]}, + ] + self.proc = NestedMultiHotProcessor() + self.proc.fit(self.samples, "codes") + + def test_vocab_matches_nested_sequence(self): + """Same vocabulary layout, so a cached vocab restores into either. + + This is what lets an existing cohort cache be reused without a rebuild. + """ + other = NestedSequenceProcessor() + other.fit(self.samples, "codes") + self.assertEqual(self.proc.code_vocab, other.code_vocab) + self.assertEqual(self.proc.code_vocab[""], 0) + self.assertEqual(self.proc.code_vocab[""], 1) + + def test_shape_is_visits_by_vocab(self): + out = self.proc.process([["A"], ["B", "C"], ["D"]]) + self.assertEqual(out.shape, (3, self.proc.vocab_size())) + self.assertEqual(out.dtype, torch.float) + + def test_marks_present_codes(self): + out = self.proc.process([["A", "C"]]) + expected = {self.proc.code_vocab["A"], self.proc.code_vocab["C"]} + self.assertEqual(set(out[0].nonzero().flatten().tolist()), expected) + + def test_repeats_collapse(self): + """A code charted five times is still one bit. + + HALO's encoder already did this (``= 1``, not ``+= 1``), so preserving + it is what keeps output identical to the index pipeline. + """ + out = self.proc.process([["A", "A", "A", "A", "A"]]) + self.assertEqual(out.max().item(), 1.0) + self.assertEqual(out[0].sum().item(), 1.0) + + def test_empty_visit_is_all_zero(self): + out = self.proc.process([["A"], [], ["B"]]) + self.assertEqual(out[1].sum().item(), 0.0) + + def test_empty_sample_is_one_empty_visit(self): + """Mirrors NestedSequenceProcessor returning a single all- row.""" + self.assertEqual(self.proc.process([]).shape, (1, self.proc.vocab_size())) + self.assertEqual(self.proc.process([]).sum().item(), 0.0) + + def test_unknown_code_maps_to_unk(self): + out = self.proc.process([["NOT_IN_VOCAB"]]) + self.assertEqual(out[0].nonzero().flatten().tolist(), + [self.proc.code_vocab[""]]) + + def test_pad_column_never_set(self): + """Column 0 must stay clear, or padding becomes indistinguishable from + a real code.""" + out = self.proc.process([["A", "B"], [], ["C"]]) + self.assertEqual(out[:, 0].sum().item(), 0.0) + + def test_none_entries_treated_as_unknown(self): + out = self.proc.process([[None]]) + self.assertEqual(out[0].nonzero().flatten().tolist(), + [self.proc.code_vocab[""]]) + + def test_vocab_edit_helpers(self): + p = NestedMultiHotProcessor() + p.fit(self.samples, "codes") + p.add({"Z"}) + self.assertIn("Z", p.tokens()) + p.retain({"A"}) + self.assertEqual(p.tokens(), {"", "", "A"}) + # Width follows the vocabulary, so it must shrink with it. + self.assertEqual(p.process([["A"]]).shape, (1, 3)) + + +if __name__ == "__main__": + unittest.main()