WorkupFlow EHR Imputation

WorkupFlow is a PyTorch implementation of a clinical workup-guided flow matching model for mixed-type EHR missing-data imputation.

Repository Contents

• workupflow/: model, preprocessing, training, sampling, evaluation, and baseline utilities.
• configs/template.yaml: template dataset configuration.
• requirements.txt: Python dependencies.

Raw datasets, private dataset configs, checkpoints, run outputs, and manuscript files are intentionally excluded from this repository.

Setup

python -m venv .venv
.venv\Scripts\activate
pip install -r requirements.txt

For Linux/macOS, activate the environment with source .venv/bin/activate.

Prepare Your Data

1. Put your local CSV under datasets/, for example datasets/my_dataset.csv.
2. Copy configs/template.yaml to configs/my_dataset.yaml.
3. Edit column names, categorical cardinalities, panels, bounds, and optional label settings.

Files under datasets/ and private configs/*.yaml are ignored by Git by default.

Train

python -m workupflow.scripts.train --dataset my_dataset --run-name my_run --epochs 10 --device auto

Evaluate

python -m workupflow.scripts.evaluate --run-dir runs/my_run --split test --nfe 30 --scenarios mcar mar mnar block mixed --holdout-rates 0.1 0.3 0.5

Sample Imputations

python -m workupflow.scripts.sample --run-dir runs/my_run --split test --n-imputations 5

Notes

• target_num contains continuous targets.
• target_cat currently expects binary categorical targets.