ECGEN

ECGEN is a research framework for generative modeling of 12-lead electrocardiograms, built on PyTorch Lightning. It provides model implementations, dataset loaders, one YAML-driven training command (ecgen train), evaluation metrics and visualization tools.

• Get started

  ---

  Install the package and train your first model on MIMIC-IV-ECG.

  Installation · Quick Start

• Models

  ---

  VAE, VQ-VAE with a PixelCNN prior, and the Pulse2Pulse GAN.

  Architecture overview

• Training

  ---

  One command for every model: ecgen train <config.yaml> [key=value ...].

  Training with ecgen train

• API reference

  ---

  Generated from the source docstrings.

  API reference

Model families

Family	Package	model.name	Status	Docs
Variational Autoencoder	ecgen.families.vae	vae	Available	VAE
VQ-VAE + PixelCNN prior	ecgen.families.vqvae	vqvae, vqvae_prior	Available	VQ-VAE
Pulse2Pulse GAN	ecgen.families.gan	pulse2pulse	Available	GAN
Diffusion	—	—	Planned	—
State Space Models	—	—	Planned	—

Datasets

Dataset	Status	Docs
MIMIC-IV-ECG (mimic_iv_ecg)	Available	MIMIC-IV-ECG
Synthetic ECG-like signals (synthetic)	Available, for smoke tests	Smoke runs
PTB-XL	Planned (config only)	—
Your own data	—	Adding a dataset

At a glance

git clone https://github.com/vlbthambawita/ECGEN.git
cd ECGEN
pip install -e .

# Train a VQ-VAE (stage 1); overrides follow the config
ecgen train configs/experiments/vqvae_mimic.yaml data.dataset.params.root=/path/to/mimic-iv-ecg

# then its prior (stage 2)
ecgen train configs/experiments/prior_mimic.yaml \
    model.params.vqvae_checkpoint=runs/vqvae_mimic/seed_42/checkpoints/last.ckpt

from ecgen.families.vae import VAELightning

model = VAELightning.load_from_checkpoint("runs/vae_mimic/seed_42/checkpoints/last.ckpt")
samples = model.sample(n_samples=16, seq_length=5000)

Note

See Training with ecgen train for the CLI, the config schema and the run directory layout. The Lightning wrappers rebuild their config from the flat hyperparameters stored in the checkpoint, so load_from_checkpoint(path) needs no extra arguments. See VAE for details.

Repository layout

ECGEN/
├── src/ecgen/       # installed package: families/, nn/, data/, training/ (runner, callbacks), cli.py
├── training/        # metrics/, losses/common.py
├── evaluation/      # metrics/, benchmark.py, visualize.py
├── visualization/   # ECG plots, latent space, training curves
├── scripts/         # shell/ (run_train_*.sh), generate/, evaluate/, data/
├── configs/         # dataset/, model/, training/, experiments/
├── single_file/     # self-contained VQ-VAE implementation
└── docs/            # this documentation

See Project Structure for details.

Building these docs locally

pip install -r requirements-docs.txt
mkdocs serve        # http://127.0.0.1:8000

Pushes to main or simple that touch the docs or the source packages rebuild the site and publish it to GitHub Pages.