[corgan-pr-prep] Overview: CorGAN PyHealth 2.0 Integration Tasks

[jalengg]

Overview

Tracking issue for the full CorGAN → PyHealth 2.0 integration on branch `corgan-pr-integration`.

CorGAN (Correlated-feature GAN) generates synthetic EHR patient records using a GAN with an autoencoder component. This integration brings CorGAN into the PyHealth 2.0 standard pattern: BaseTask task function, BaseModel inheritance, updated examples, dead-code removal, Google-style docstrings, and a MECE integration test suite.

Tasks

#	Issue	Status
T1	#25 Create corgan-pr-integration branch from upstream/master	✅ Done
T2	#26 Create CorGANGenerationMIMIC3 task function	✅ Done
T3	#28 Refactor CorGAN to full PyHealth 2.0 BaseModel compliance	🔄 In progress
T4	#30 Update CorGAN training example	Pending
T5	#27 Update CorGAN generation example	Pending
T6	#29 Remove bespoke CorGAN dataset class if any	Pending
T7	#31 Update CorGAN docstrings to Google/PyHealth style	Pending
T8	#32 Write CorGAN integration tests	Pending
T9	#33 Final verification and push	Pending

Methodology

Each task follows two-stage review: spec compliance first, then code quality. Both must pass before the task is closed.

Branch: `corgan-pr-integration`
Fork: `jalengg/PyHealth`
Upstream PR target: `sunlabuiuc/PyHealth`

[jalengg]

Environment Note: Extra Dependencies from upstream/master

During T3 smoke testing we hit two ModuleNotFoundError crashes that are not bugs in the CorGAN integration — they come from pyhealth/datasets/base_dataset.py on upstream/master, which has accumulated heavy new imports over the past year that were not present when the original HALO integration was developed.

What happened

Running a smoke test against the corgan-pr-integration worktree (based on upstream/master at commit 436bba7e) failed with:

ModuleNotFoundError: No module named 'pyarrow'
ModuleNotFoundError: No module named 'dask'

Both are required by base_dataset.py before any model or dataset class can be imported.

When and why these were added upstream

Package(s)	Merged	Reason
pyarrow, dask, dask.distributed	May 2025	Parallelised the patient-level task transformation loop using Dask's distributed scheduler. pyarrow was added alongside it for Parquet-based intermediate storage during parallel writes. The change achieved a 2–3× wall-time speedup on MIMIC-III set_task() calls.
litdata, litdata.streaming, litdata.processing	Dec 2025	Added Lightning's litdata streaming library to support caching set_task() results to disk. Previously, calling set_task() on a large dataset always re-ran the full transformation. With litdata, the result is written once as a BinaryWriter stream and read back on subsequent calls.

Both packages are imported unconditionally at module load time in base_dataset.py lines 17–32, so any from pyhealth.datasets import ... will fail if they are absent from the environment.

Fix

Install in the project venv before running any tests or smoke tests on this branch:

/u/jalenj4/pehr_scratch/venv/bin/pip install pyarrow "dask[distributed]" litdata

These were installed on Feb 25, 2026 during T3 setup. The venv is ready for T8 and T9. If you recreate the venv from scratch, re-run the above.

Why requirements.txt does not list them

The upstream requirements.txt has not been updated to reflect these additions — it still lists the original minimal dependencies (torch, sklearn, polars, etc.). This is a pre-existing gap in the upstream repo, not something introduced by this integration. The requirements.txt-only install path is currently broken for any PyHealth 2.0 code that touches BaseDataset. This is worth flagging in the upstream PR review.