Add NestedMultiHotProcessor and vectorise HALO's visit encoding

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

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:

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)

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:

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 <pad> (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, "<pad>", "<unk>")
+                ]
+                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}")

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, "<pad>", "<unk>"):
-                        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)}"