[margaretdu]

Summary

GAN was an image-only nn.Module from before PyHealth 2.0, with no tests and no way to train on EHR data. This PR adds a code-set mode so it can train on MIMIC-III, plus unit tests and an example notebook.

• pyhealth/models/gan.py: GAN is now a BaseModel with two modes.
  • Code-set mode (pass a dataset whose single input feature is multi_hot, e.g. EHRCodeSetGenerationMIMIC3): MLP generator and discriminator that generate synthetic patients as sets of codes.
  • Image mode (pass input_channel and input_size): the existing CNN networks, unchanged. An unsupported input_size now raises a ValueError instead of building a model with no networks.
  • Adds train_model() for the adversarial loop (the Trainer does not apply to GANs, as with the other generators), generate() to decode patients in MedGAN.generate()'s output format, and a loss-only forward() for a batch.
• tests/core/test_gan.py: 14 tests on pseudo data, following test_mlp.py: initialization, forward pass, backward pass (each loss reaches only its own network), training, generation, input validation, and image mode.
• examples/gan_mimic3.ipynb: trains on Synthetic MIMIC-III diagnoses in about 30 seconds on CPU, plots both losses, generates synthetic patients and compares them with held-out real patients.
• docs/api/models/pyhealth.models.GAN.rst: describes both modes and links both example notebooks.

Notes for review

• Results are modest by design. This is a minimal example: synthetic patients average 19.5 codes vs. 16.7 for real ones, but a few common codes are over-produced (e.g. 4019 in 99% of synthetic vs. 49% of real patients). The notebook states this and points to MedGAN for a purpose-built generator.
• Training length. Longer runs (100 epochs) led the discriminator to win outright and the generator to collapse, so the notebook trains for 30 epochs.
• Saved outputs. The notebook was executed headlessly; progress-bar fragments were merged as Jupyter displays them, and a LitData "newer version available" warning (which quotes a local install path) was removed from the saved output. Nothing else was edited.
• Positional image-mode calls change. dataset is now the first argument, so GAN(3, 128) must become GAN(input_channel=3, input_size=128). The existing chest X-ray notebook already uses keywords.

Testing

• python -m unittest discover -t tests -s tests/core -p "test_gan.py": 14 passed.
• Docstring example runs as a doctest.
• tools/check_pr_rules.py passes on the branch.
• Notebook executed end to end on Python 3.13 (CPU).

🤖 Generated with Claude Code

[jhnwu3]

/review

[github-actions]

Test Coverage

Total: 68.07% (16790 / 24666 lines)

Modified file	Coverage	Lines
pyhealth/models/gan.py	95.76%	113 / 118

Workflow run · c75ceaa

[github-actions]

AI review: findings

The PR turns the image-only GAN into a BaseModel with a code-set mode: an MLP GAN over a multi_hot feature (e.g. EHRCodeSetGenerationMIMIC3), plus train_model(), generate() and a loss-only forward(), while the CNN image mode is kept unchanged. It is a distinct method from MedGAN and CorGAN (a vanilla GAN, not autoencoder-based), so it is not redundant as a component. Blocking problem (medium), pyhealth/models/gan.py:75: the constructor's first positional parameter is now dataset, so old calls such as GAN(3, 128) pass 3 as the dataset and fail with an opaque AttributeError: 'int' object has no attribute 'input_schema' raised from BaseModel.init. The in-repo caller already uses keywords, but the break is only noted in the PR text, not in code or docs. Non-blocking (low), pyhealth/models/gan.py:378: generate() and the _idx_to_code setup copy the decode logic of MedGAN.generate (pyhealth/models/generators/medgan.py), which a shared helper would avoid. Tests: 14 new unit tests in tests/core/test_gan.py cover initialization, forward, per-network backward isolation, training, generation, validation and image mode, at about the level of test_medgan.py, and there is a doctest. I checked the notebook's parameter count (791,495) against the architecture and it matches; I did not execute the tests or the notebook. Nothing else in the repo enumerates models or imports GAN. Overall risk is low to moderate, with the constructor break as the main concern for existing image-mode users.

_{Reviewed commit 2f3b0c3}

[github-actions]

[medium] Positional constructor change turns old GAN(3, 128) calls into an opaque AttributeError

The original signature was (input_channel, input_size, hidden_dim). The new first positional parameter is dataset, so GAN(3, 128) passes 3 as the dataset and BaseModel.init fails with AttributeError: 'int' object has no attribute 'input_schema'. The break is documented only in the PR description. Fix: keep input_channel and input_size as the first positional parameters and make dataset keyword-only after them, e.g. def __init__(self, input_channel=None, input_size=None, hidden_dim=128, *, dataset=None, **kwargs), so existing calls keep working. If the project records API changes in CHANGELOG.md, note the change there.

[margaretdu]

(input_channel, input_size, hidden_dim) is an artefact of Pyhealth 1.0, I'm not sure if anyone would be trying this/if there's already broader handling for such versioning issues.

[github-actions]

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change reorders the public GAN constructor parameters (dataset is now first) and changes GAN's base class from nn.Module to BaseModel, which are API contract changes. The effect is limited to direct users of GAN, since nothing else in the repo imports it, but the rule treats reordered parameters as core.

_{Reviewed commit 2f3b0c3}