pyvae

Informed Variational Autoencoder (iVAE) for gene expression, in PyTorch.

pyvae implements a VAE whose encoder is biologically constrained: each gene can only influence the pathways that contain it, according to the Reactome database. The result is a latent space whose dimensions correspond to interpretable biological pathways rather than opaque features.

Why an informed VAE?

A standard Variational Autoencoder learns a compressed latent representation $z$ of the input $x$ by maximising the Evidence Lower Bound (ELBO):

$$ \text{ELBO} = \mathbb{E}_{q(z|x)}[\log p(x|z)] - \text{KL}(q(z|x),|,p(z)) $$

Gene expression is high-dimensional (~20 000 genes) but structured: genes act together in pathways. pyvae restricts the encoder's first layer so only the gene → pathway connections that exist in Reactome are allowed. A binary mask enforces this on the weight matrix: $W_{\text{eff}} = W \odot M$, where $M_{ji}=1$ iff gene $i$ belongs to pathway $j$. This reduces parameters, injects prior biological knowledge, and makes each latent node interpretable as a pathway.

Validation benchmark

The canonical end-to-end validation of pyvae lives in experiments/kang/: a Snakemake grid of 28 training jobs on the Kang 2018 PBMC interferon-β dataset, plus a showcase notebook that examines one model closely with saturation, unit fidelity, decoder-side pathway ranking, gene attribution, and counterfactual prediction. The sweep exercises the architecture ablation, the encoder-fidelity constraints, and the sample-size behaviour against a model-free baseline. See experiments/kang/README.md for how to run it and experiments/kang/FINDINGS.md for what the results show.

Features

Models

InformedLinear is a masked linear layer that enforces arbitrary connectivity priors. InformedVAE builds a full encoder/decoder around it, using the reparameterisation trick and a combined reconstruction, KL, and L2 loss. You configure it through constructor kwargs. likelihood picks between "gaussian" (MSE on log1p-normalized expression, the default) and "nb" (a negative-binomial likelihood on raw counts, implemented scVI-style in pure PyTorch with no scvi-tools dependency). n_cov sets the width of an auxiliary one-hot covariate (cell type, condition, or both), which the model concatenates into the encoder and decoder to support conditional inference.

The encoder has four optional constraints: init="fan_in" for per-unit Xavier scaling, normalize="batch" for BatchNorm on the pre-activation (which keeps units off tanh's flat tail), nonneg_encoder=True for a softplus reparameterisation that removes sign freedom (so a unit cannot encode its pathway upside-down), and standardize_input=True for BatchNorm on the input (which removes scale freedom). informed_decoder=True applies the same Reactome mask symmetrically on the decoder, which is what enables faithful counterfactual prediction.

Training

train_ivae is the legacy training loop: gradient clipping plus early stopping. train_ivae_modern trains the count model instead, adding KL warmup, AdamW with decoupled weight decay, cosine LR annealing, and best-weight restore.

train_ivae_modern also accepts cov_train and cov_val to thread one-hot covariates through both encoder and decoder — required for the counterfactual read-out and recommended whenever cell composition varies between conditions. Datasets are moved to the training device once at the start of training, so batches cost a device-side gather rather than a host-to-device copy per epoch; a full-panel Kang run takes about 4 minutes on a single A100.

Interpretation

InformedVAE.predict_counterfactual (NB models only) encodes a cell under one covariate and decodes it under another, returning predicted counts. It uses the posterior mean rather than a stochastic sample, so the prediction stays deterministic, which is what you want for a question like "what would this control cell look like if stimulated?" bayes_factor_da computes a signed Bayes factor per pathway between two groups of cells, based on Monte-Carlo pairing of raw pathway activations. integrated_gradients attributes a single pathway's activation back to each input gene (Sundararajan et al., 2017), implemented in pure PyTorch.

Four pathway-level read-outs cover the gap between a VAE and a traditional pathway-ranking workflow. pathway_unit_fidelity correlates each unit's activation against the mean z-scored expression of its own member genes, so you can audit whether a unit actually represents the pathway named after it (and whether its sign is correct). differential_expression samples posterior per-gene differential expression between two groups (Boyeau et al. 2019), with observed-detection filtering so decoder-confident but barely-observed phantom genes do not float to the top of the ranking. pathway_activity runs a tie-corrected competitive Mann-Whitney rank test of a pathway's member genes against the rest of the panel on the output of differential_expression, which is the ranking most biologists actually want; it also reports top_gene_share so you can tell when a hit is really one gene wearing a pathway's name. pseudobulk_paired_test collapses cells to donor-level pseudobulk and runs a paired Wilcoxon per gene — the Squair et al. 2021 correction for cell-level pseudo-replication, which inflates cell-level p-values badly when donors are the real unit of replication.

Data helpers

Reactome helpers (build_model_config, sync_gexp_adj) and a Kang PBMC dataset loader (load_kang) cover end-to-end pipelines. swap_condition flips the one-hot condition columns in a covariate DataFrame, which is how you build the cov_to argument for predict_counterfactual. set_all_seeds seeds Python's random, NumPy, and PyTorch (CPU and CUDA) in a single call, for reproducible runs.

Installation

pip install pyvae                      # from PyPI
conda install -c conda-forge pyvae     # from conda-forge
pixi add pyvae                         # into a pixi project

pyvae is a pure-Python, noarch package and runs on Linux, macOS, and Windows.

GPU acceleration

pyvae only requires torch>=2.3,<3 and never pins a platform. Which hardware you can accelerate depends entirely on which PyTorch build your install command pulls. Use this table to pick the right path:

CUDA and Apple-Silicon users need no extra steps. Verify at runtime with python -c "import torch; print(torch.cuda.is_available())" (CUDA) or torch.backends.mps.is_available() (macOS).

AMD ROCm users must install torch from PyTorch's ROCm index. ROCm wheels exist neither on PyPI nor on conda-forge, so conda install is a dead end for AMD GPUs. Install ROCm torch first, then pyvae; pip then sees the constraint already satisfied and leaves torch alone:

pip install torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/rocm7.2   # match your ROCm runtime
pip install pyvae

To pin an explicit CUDA stack via conda instead of relying on auto-detection:

conda create -n pyvae-cuda12 -c conda-forge -c pytorch -c nvidia \
  pyvae pytorch pytorch-cuda=12.*

Match the CUDA/ROCm version to your driver/runtime; the official PyTorch selector gives exact wheels.

Quickstart

import pandas as pd
import torch
import pyvae

# 1. Load and preprocess data (Kang PBMC)
adata = pyvae.load_kang(data_folder="./data", n_genes=3000)

# 2. Train/validation split
train_mask = adata.obs["split"] == "train"
x_train = pd.DataFrame(adata[train_mask].X.toarray(), columns=adata.var_names)
x_val   = pd.DataFrame(adata[~train_mask].X.toarray(), columns=adata.var_names)

# 3. Build the biological (Reactome) configuration and adjacency
config = pyvae.build_model_config(
    genes=x_train.columns,
    model_kind="ivae_reactome_hierarchical_d2",
    resources_dir="./resources",
)
adj = torch.tensor(config.model_layer[0].values, dtype=torch.float32)

# 4. Build and train the model
model = pyvae.InformedVAE(adj=adj, latent_dim=64, l2_lambda=1e-5)
model, history = pyvae.train_ivae(
    model, x_train, x_val,
    epochs=500, batch_size=32, patience=50, lr=1e-4, device="cpu",
)

# 5. Encode new cells — `h` is the interpretable pathway activation
z_mean, z_log_var, h = model.encode(torch.tensor(x_val.values, dtype=torch.float32))

Development

This repo is managed with pixi. The manifest defines several environments (default, ci, cuda12, rocm72, publish):

pixi install              # create the default environment
pixi run pytest tests/ -v # run the test suite
pixi run fixcode          # ruff lint + import sort
pixi run formatcode       # ruff format

Without pixi, a standard editable install works too:

pip install -e ".[dev]" && pytest tests/ -v

Releasing

This repo publishes to both PyPI and conda-forge. Before releasing, bump the version in pyproject.toml (both [project].version and [tool.pixi.package].version).

PyPI (Trusted Publishing in CI on a v* tag, or locally via the isolated publish env):

pixi run -e publish build-dist     # build sdist + wheel into dist/
pixi run -e publish check-dist     # twine metadata check
pixi run -e publish publish-pypi   # upload (set TWINE_REPOSITORY=testpypi for TestPyPI)

conda (built as a single noarch package by pixi):

pixi publish --target-dir ./conda-dist          # build + copy locally (no upload)
pixi publish --to https://prefix.dev/<channel>  # or push to a channel directly

The pyvae-feedstock handles conda-forge distribution instead of a direct upload, and the package is already live there. regro-cf-autotick-bot picks up new PyPI releases automatically and opens a version-bump PR on the feedstock, so no action is needed here. The original submission recipe and instructions for manual recipe changes live in conda-recipe/. On a tagged release, .github/workflows/release.yml builds and verifies both artifacts and attaches the noarch conda package to the GitHub Release.

Project layout

pyvae/
├── __init__.py    public API
├── utils.py       seeding utilities
├── layers.py      InformedLinear (masked layer)
├── components.py  Encoder, DenseDecoder, CountDecoder, likelihoods, helpers
├── models.py      InformedVAE (composition of components)
├── train.py       training loops: train_ivae, train_ivae_modern
├── datasets.py    Kang dataset loader
├── bio.py         Reactome adjacency helpers, swap_condition
├── interpret.py   bayes_factor_da, integrated_gradients,
│                   pathway_unit_fidelity, differential_expression,
│                   pathway_activity, pseudobulk_paired_test
tests/             contract + unit tests
conda-recipe/      conda-forge recipe + submission guide
experiments/kang/  canonical validation benchmark (Kang 2018 IFN-β PBMCs);
                   see experiments/kang/README.md and FINDINGS.md

Acknowledgements

pyvae is a PyTorch reimplementation that builds on earlier Keras informed-VAE work directed by Carlos Loucera, with contributions from Pelin Gundogdu (babelomics/ivae_scorer), Alberto Esteban-Medina (albertoem77/robustness_informed_TFM), and Sara Fernandez (saraafdezz/robustness_informed_TFG).

License

MIT © Amir Aynede, Carlos Loucera