Coding Journal Club: Running ViVAE on Tabula Sapiens Bone Marrow
TL;DR: ViVAE is a variational-autoencoder embedding from the Saeys Lab, built to keep global structure that t-SNE and UMAP tend to break. We ran the official scRNA-seq tutorial end to end on 12,297 bone marrow cells. ViVAE had the best global structure score (0.569 vs. 0.419 for t-SNE and 0.410 for UMAP) but the worst local score (0.325 vs. 0.432 and 0.345). Useful, but not a free upgrade over UMAP.
Why ViVAE
t-SNE and UMAP are neighbour-embedding methods: they are good at keeping each cell next to its nearest neighbours, and bad at keeping distances between far-apart groups meaningful. Cluster-to-cluster distances on a UMAP plot mostly don't mean much.
ViVAE (paper: Interpretable models for scRNA-seq data embedding with multi-scale structure preservation) trains a variational autoencoder with an added MDS-style loss that pushes the 2D layout to respect distances at every scale. According to the README, it:
- aims for multi-scale structure preservation, which matters most for data with trajectories, outlier populations or suspected batch effects
- ships encoder indicatrices, which show where and in which direction the embedding is locally stretched
- is parametric, so new cells can be embedded without retraining
- integrates with FlowSOM
Its companion package ViScore scores embeddings without labels (Local/Global structure preservation) and with labels (xNPE per population).
Everything below follows the repo's tutorials/example_scrnaseq.ipynb. We ran every step below (print statements tidied for readability), and every number shown comes from our own run, not from the paper.
Test setup
- Machine: Linux workstation, CPU only, 16 threads
- Versions: Python 3.11.9, vivae 0.0.2, viscore 0.0.2, flowsom 0.2.3, torch 2.14.0, scanpy 1.11.5, scikit-learn 1.9.1, umap-learn 0.5.9.post2
- Data: Tabula Sapiens bone marrow from CELLxGENE, 12,297 cells, 15 annotated cell types (274 MB
.h5ad) - Cost: about 3.5 minutes for the whole pipeline; peak memory 10.4 GB
Step 1: Installation
conda create --name ViVAE python=3.11.9
conda activate ViVAE
pip install numpy==1.26.3 pandas==2.2.0 matplotlib==3.8.2 scipy==1.12.0 \
scikit-learn==1.4.0 scanpy==1.9.8 pynndescent==0.5.11 \
umap-learn==0.5.6 pyemd==1.0.0
pip install git+https://github.com/saeyslab/ViVAE.git \
git+https://github.com/saeyslab/ViScore.git \
git+https://github.com/saeyslab/FlowSOM_Python.git
# Needed in practice, see below
pip install "setuptools<81" umap-learn==0.5.9.post2
Two things broke with the pinned versions from the README, so we added that last line:
import viscorefailed withModuleNotFoundError: No module named 'pkg_resources'. ViScore usespyemd, which still importspkg_resources, and new setuptools releases no longer ship it. Installingsetuptools<81fixes it.- Installing FlowSOM upgraded scikit-learn (to 1.9.1 here), after which umap-learn 0.5.6 crashed with
check_array() got an unexpected keyword argument 'force_all_finite'. umap-learn 0.5.9.post2 works. If you skip FlowSOM, the original pins should be fine.
Backend and reproducibility. ViVAE reads three environment variables, and they must be set before import vivae:
import os
os.environ['VIVAE_DETERMINISTIC'] = '1' # default: deterministic PyTorch ops (slower)
os.environ['VIVAE_CUDA'] = '0' # default '1': use CUDA if available
os.environ['VIVAE_MPS'] = '0' # default '0': Apple GPU off
import vivae as vv
Determinism is on by default, but the package switches it off when it runs on CUDA. The README recommends the CPU backend for reproducing the tutorial results, which is what we used.
Step 2: Download and preprocess
wget -O scrnaseq.h5ad https://datasets.cellxgene.cziscience.com/f56e0393-c2f3-42f3-93f1-7c8a4ee49d75.h5ad
import scanpy as sc
import numpy as np
hd = sc.read_h5ad('scrnaseq.h5ad')
labs = np.array(hd.obs['cell_type'])
sc.pp.scale(hd, max_value=10.) # z-standardisation
sc.tl.pca(hd, svd_solver='arpack', n_comps=100)
pc = hd.obsm['X_pca']
This file is already normalised, so there is no normalize_total/log1p. With a raw count matrix you would run those first. The tutorial keeps 100 PCs rather than the usual 50 as the more conservative choice.
Step 3: k-NN graph and denoising
knn = vv.make_knn(x=pc, k=1000, fname='scrnaseq_knn.npy') # saved, reloaded on the next run
pc_denoised = vv.smooth(x=pc, knn=knn, k=50, coef=1., n_iter=1)
smooth moves each point a fraction coef of the way toward the mean of its k nearest neighbours. With coef=1. each cell is replaced by the mean of its 50 neighbours. The tutorial does this to sharpen the layout's skeleton.
Building the 1,000-neighbour graph took 53 s. The smoothing itself took under a second.
Step 4: Train ViVAE
import torch
np.random.seed(42)
torch.manual_seed(42)
model = vv.ViVAE(input_dim=pc.shape[1], latent_dim=2, random_state=42)
model.fit(X=pc_denoised, n_epochs=100, batch_size=512, verbose=True)
emb = model.transform(pc_denoised)
Training took 40–52 s on CPU. With verbose=True it prints one line per epoch, with each loss term reported separately:
Epoch 1/100 recon: 73.6558 kldiv: 0.7139 geom: 0.0000 egeom: 0.0000 mds: 431.6573 imit: 0.0000
Epoch 2/100 recon: 36.1623 kldiv: 0.3566 geom: 0.0000 egeom: 0.0000 mds: 192.0285 imit: 0.0000
...
Epoch 100/100 recon: 0.0802 kldiv: 0.0259 geom: 0.0000 egeom: 0.0000 mds: 0.8671 imit: 0.0000
With default weights, three terms are active: reconstruction (lam_recon=1), KL divergence to the latent prior (lam_kldiv=1) and the MDS distance loss (lam_mds=100, the value the docstring recommends for scRNA-seq; it suggests 10 for cytometry). The geometric and imitation terms are off by default.
Defaults worth knowing, since the tutorial overrides two of them: n_epochs=50, batch_size=256, learning_rate=1e-3. The tutorial trains for 100 epochs with batch size 512.
Step 5: t-SNE and UMAP for comparison
from sklearn.manifold import TSNE
import umap
emb_tsne = TSNE(n_components=2, perplexity=30., random_state=42).fit_transform(pc)
model_umap = umap.UMAP(random_state=42)
model_umap.fit(pc)
emb_umap = model_umap.transform(pc)
Like the tutorial, we give t-SNE and UMAP the plain PCs with default settings. Each took well under a minute here (about 20 s for t-SNE, about 30 s for UMAP). Setting random_state makes UMAP single-threaded.
vv.plot_embedding(emb, labels=labs) draws any of the three embeddings coloured by cell type.
Step 6: Encoder indicatrices
ei = model.encoder_indicatrices(X=pc_denoised, radius=1e-3, n_steps=30, n_polygon=200)
fig, ax = vv.plot_indicatrices(indicatrices=ei, scale_factor=4e2, show=False)
encoder_indicatrices returns an EncoderIndicatome object, not a number per cell. It samples a 30 × 30 grid over the latent space and, at each grid point, maps a tiny sphere from the input space into 2D. Each resulting polygon shows how much, and in which direction, the embedding stretches or squeezes the data at that spot. The polygons are drawn on top of the embedding, so treat them as a distortion map, not a colour scale. Using it as a per-cell colour, e.g. plt.scatter(c=ei), fails with a TypeError.
In the tutorial's reading of this plot, the region holding monocytes and granulocytes is artificially shrunk and the plasma cells show a "run-away" stretch. The plasma cell stretch also shows up in the xNPE scores in Step 8.
Step 7: Unsupervised scores with ViScore
import viscore as vs
for name, e in [('t-SNE', emb_tsne), ('UMAP', emb_umap), ('ViVAE', emb)]:
s = vs.score(hd=pc, ld=e) # returns a dict: Sl, Sg, RNX, QNX
print(f"{name:<6} Local SP {s['Sl']:.3f} Global SP {s['Sg']:.3f}")
t-SNE Local SP 0.432 Global SP 0.419
UMAP Local SP 0.345 Global SP 0.410
ViVAE Local SP 0.325 Global SP 0.569
Both scores are areas under an RNX neighbourhood-preservation curve, bounded by −1 and 1, where 0 means no better than a random embedding. Local SP weights small neighbourhoods more (log scale); Global SP weights all neighbourhood sizes equally. The first vs.score call took about 20 s because of numba compilation; later calls take under a second.
What this says: ViVAE clearly wins on global structure and comes last on local structure. t-SNE keeps immediate neighbourhoods best. Your numbers will move a little with hardware and library versions, but that pattern matches what the tutorial reports.
Always score against the undenoised pc. ViVAE was trained on smoothed data, so scoring it against pc_denoised flatters it: its Global SP jumps from 0.569 to 0.687 and Local SP from 0.325 to 0.468. The tutorial warns about exactly this.
Step 8: Which populations are distorted (xNPE)
xNPE (Extended Neighbourhood-Proportion-Error) uses the labels. It compares how mixed each population's neighbourhoods are in the original data versus the embedding. It gives one score per population, where 0 is perfect and, with the default baseline correction, 1 is random-level error.
xnpe_umap = vs.xnpe(hd=pc, ld=emb_umap, annot=labs, knn=knn) # dict: population -> score
xnpe_vivae = vs.xnpe(hd=pc, ld=emb, annot=labs, knn=knn)
for pop in sorted(xnpe_vivae, key=xnpe_vivae.get, reverse=True)[:4]:
print(f"{pop:<33} {xnpe_vivae[pop]:.3f}")
print("UMAP better than ViVAE for:", [p for p in xnpe_vivae if xnpe_umap[p] < xnpe_vivae[p]])
plasmablast 0.600
mature NK T cell 0.363
plasma cell 0.276
CD8-positive, alpha-beta T cell 0.228
UMAP better than ViVAE for: ['CD8-positive, alpha-beta T cell', 'erythrocyte', 'mature NK T cell', 'memory B cell', 'monocyte', 'neutrophil']
Pass the knn graph from Step 3. Without it (or an explicit k) the function errors out. The two calls took about a minute together. vs.plot_xnpe_barplot turns the results into a bar chart, and vs.neighbourhood_composition_plot shows which neighbours a badly embedded population picked up.
The plasma-lineage cells (plasmablasts and plasma cells) are ViVAE's worst populations, which matches the stretching seen in the indicatrices. UMAP beats ViVAE on 6 of the 15 populations, including the T/NK populations the tutorial also flags.
Bonus: FlowSOM on top of the embedding
FlowSOM wants an AnnData object with channel/marker names, so wrap the PCs first:
import anndata as ad
import pandas as pd
import flowsom as fs
hdf = ad.AnnData(pc)
cols = [f'PC{i+1}' for i in range(pc.shape[1])]
hdf.var = pd.DataFrame({'n': range(pc.shape[1]), 'channel': cols, 'marker': cols})
fsom = fs.FlowSOM(hdf, n_clusters=40, xdim=12, ydim=12, seed=42)
fig, ax = vv.plot_embedding(embedding=emb, labels=labs, s=1., fsom=fsom,
dr_model=model, fsom_show_nodes=False, show=False)
Clustering happens in PC space, not on the embedding. dr_model=model lets ViVAE project the FlowSOM nodes and minimum spanning tree into its 2D layout, which works because the model is parametric. The same call with dr_model=model_umap does it for UMAP, which also has .transform; t-SNE can't. The whole step took about 4 s.
Takeaways
Reach for ViVAE when:
- distances between populations matter: lineages, trajectories, how far apart two groups really are (best Global SP here by a wide margin)
- you want to see where the 2D layout distorts the data (indicatrices, xNPE)
- you need to embed new cells into an existing layout
Stay with t-SNE/UMAP when:
- you mainly need tight local neighbourhoods (both beat ViVAE on Local SP here)
- you want a simpler stack. The ViVAE pipeline needs a 1,000-neighbour graph, PyTorch and a few version pins to install cleanly, and our full run peaked at 10.4 GB of memory for 12k cells.
Whichever you choose: score embeddings against the unsmoothed input, and check per-population errors before trusting what a 2D plot says about a specific cell type.
Links
- Paper: Interpretable models for scRNA-seq data embedding with multi-scale structure preservation (bioRxiv)
- Code: ViVAE (tutorial:
tutorials/example_scrnaseq.ipynb) · ViScore · FlowSOM_Python - Benchmarking: ViScore benchmarking framework for comparing DR methods across datasets on an HPC cluster
Part of the Coding Journal Club series, where we run published bioinformatics tools on real data and report what actually happens.