Coding Journal Club: Running ViVAE on Tabula Sapiens Bone Marrow

· 9 min read · Tutorials

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:

  1. import viscore failed with ModuleNotFoundError: No module named 'pkg_resources'. ViScore uses pyemd, which still imports pkg_resources, and new setuptools releases no longer ship it. Installing setuptools<81 fixes it.
  2. 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

Part of the Coding Journal Club series, where we run published bioinformatics tools on real data and report what actually happens.

single-cell dimensionality-reduction machine-learning python coding-journal-club vivae umap