[medgan-pr-prep] T4: Update training example (medgan_mimic3_training.py)

[jalengg]

Goal

Create/rewrite examples/medgan_mimic3_training.py as a clean PyHealth 2.0 example using the new API.

Template

from pyhealth.datasets import MIMIC3Dataset, split_by_patient
from pyhealth.tasks import medgan_generation_mimic3_fn
from pyhealth.models import MedGAN

root = "/srv/local/data/physionet.org/files/mimiciii/1.4"

dataset = MIMIC3Dataset(root=root, tables=["diagnoses_icd", "admissions"])
sample_dataset = dataset.set_task(medgan_generation_mimic3_fn)
print(f"Total samples: {len(sample_dataset)}")

train, val, test = split_by_patient(sample_dataset, [0.8, 0.1, 0.1])

model = MedGAN(
    dataset=sample_dataset,
    latent_dim=128,
    hidden_dim=128,
    autoencoder_hidden_dim=128,
    discriminator_hidden_dim=256,
    minibatch_averaging=True,
    save_dir="./medgan_save/",
)

model.train_model(
    train_dataset=train,
    val_dataset=val,
    ae_epochs=100,
    gan_epochs=100,
    batch_size=32,
)

synthetic = model.synthesize_dataset(num_samples=1000)
print(f"Generated {len(synthetic)} patients")
print("Sample:", synthetic[0])

Acceptance criteria

• File is examples/medgan_mimic3_training.py
• Uses only PyHealth 2.0 API (MIMIC3Dataset, set_task, MedGAN, train_model, synthesize_dataset)
• No raw CSV/numpy imports
• MIMIC-III path is a variable, not hardcoded in function calls

[jalengg]

Implementation notes

Port repo has two example files

The port repo (PyHealth-Medgan-Corgan-Port) has:

• examples/generate_synthetic_mimic3_medgan.py — generation-focused (load model, generate, save CSV)
• examples/synthetic_data_generation_mimic3_medgan.py — full pipeline (train + generate)

The T4 training example should be a clean rewrite, not a patch of either. The old files use raw CSV loading and numpy binary matrices — none of that belongs in PyHealth 2.0.

What the new training example must NOT do

• No pd.read_csv() or raw file loading
• No numpy binary matrix construction
• No from_binary_matrix() or from_count_matrix() (removed in T3)
• No pretrain_autoencoder() called directly by user (folded into train_model())
• No manual DataLoader construction

Required imports

from pyhealth.datasets import MIMIC3Dataset, split_by_patient
from pyhealth.tasks import medgan_generation_mimic3_fn
from pyhealth.models import MedGAN

Note: MedGAN must be exported from pyhealth/models/__init__.py — check if T3 adds this export or if T4 needs to verify it.

MIMIC-III tables needed

tables=["diagnoses_icd", "admissions"] — both are required since the task joins on hadm_id from admissions.

Two-phase training exposed to user via single call

model.train_model(
    train_dataset=train,
    val_dataset=val,      # optional
    ae_epochs=100,
    gan_epochs=100,
    batch_size=32,
    ae_lr=0.001,
    gan_lr=0.0002,
)

The user does not call pretrain_autoencoder separately.