01 Quickstart: train a Nicheverse codebook#

Nicheverse learns two coupled vector-quantized codebooks from imaging spatial transcriptomics: a cell codebook (recurrent transcriptional states) and a neighborhood codebook (recurrent multicellular niches), coupled by cross-attention.

This notebook trains an end-to-end model on a real MERFISH mouse retina dataset (Vizgen, 4 samples), then reads back the codes and summarizes them.

The three inputs Nicheverse needs from an AnnData:

  1. raw counts in adata.X

  2. micron coordinates in adata.obsm['spatial']

  3. a sample column adata.obs['sample_id'] (the neighbor graph is built within each sample only)

Demo note: we use ~25 epochs so this runs in a few minutes; a real cohort run uses num_epochs ~300. The recommended default encoder is mlp_deep (a SwiGLU pre-norm residual MLP that gives the healthiest raw codebook on sparse imaging counts).

import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import scanpy as sc
import nicheverse as nv
from nicheverse import ModelConfig, TrainConfig
# bundled example data lives next to the notebooks, in ../examples/data
DATA = os.path.join('..', 'examples', 'data')
print('nicheverse', nv.__version__)

nicheverse 0.2.0

Load the data#

We use the bundled real MERFISH mouse retina dataset (Vizgen, 4 samples). read_spatial standardizes any AnnData or .h5ad path: it guarantees obsm['spatial'] and the sample column exist. adata.X here is raw counts; train_model applies normalize + log1p internally.

adata = nv.read_spatial(os.path.join(DATA, 'merfish_retina.h5ad'),
                        sample_col='sample_id', spatial_key='spatial')
print(adata)
print('samples:', list(adata.obs["sample_id"].astype(str).unique()))
print('cells:', adata.n_obs, ' genes:', adata.n_vars)

AnnData object with n_obs × n_vars = 113385 × 368
    obs: 'sample_id'
    obsm: 'spatial'
samples: ['VZG105a_WT1', 'VZG105a_WT2', 'VZG105a_WT3', 'VZG105a_WT4']
cells: 113385  genes: 368

Configure and train#

ModelConfig describes the architecture (encoder, codebook sizes). TrainConfig describes the optimization and the spatial neighbor graph.

We pass a fixed batch_size=2048 here. TrainConfig(batch_size='auto') is also supported: it resolves the batch from the panel size and GPU memory and scales the learning rate by sqrt(effective_batch / 2048), which is useful for throughput on large panels. Note that a very large batch gives the per-batch codebook diversity term fewer updates, so for the evenest codebook on a dataset this size a moderate fixed batch is a good choice; the resolved value is always recorded in training_runtime.json. These fields are the reference defaults (spatial_graph='knn_radius' at radius=50 um, k_neighbors=20, neighborhood_aggregation='weighted_mean', encoder_type='mlp_deep', quantizer_type='vq', seed=9, learning_rate=3e-4); only num_epochs (300 -> 30) and batch_size (32768 -> 2048) are shrunk here for a fast demo.

mc = ModelConfig(
    input_dim=adata.n_vars,
    gene_names=tuple(adata.var_names.astype(str)),
    encoder_type='mlp_deep',       # library default encoder
    quantizer_type='vq',           # recommended default quantizer
    cell_num_embeddings=256,
    neighborhood_num_embeddings=32,
)
tc = TrainConfig(
    num_epochs=30,                 # default is 300; shrunk for a fast demo
    batch_size=2048,               # default is 32768; shrunk for the small demo panel ('auto' also works)
    k_neighbors=20,
    spatial_graph='knn_radius', radius=50.0,   # reference defaults
    neighborhood_aggregation='weighted_mean', # reference default
    save_best=False, seed=9,                  # default seed
)
CKPT = 'runs/quickstart'
model, adata = nv.train_model(adata, CKPT, model_config=mc, train_config=tc, sample_col='sample_id')

[nicheverse] epoch 1/30 | total=1.6397 cell=0.9279 neigh=0.7118 | perp c/n=171.5/24.9 | active c/n=256/32 | gini c/n=0.23/0.16 | lr=3.00e-04 gnorm=6.16 | 2.2s 51432 cells/s

[nicheverse] epoch 2/30 | total=0.8889 cell=0.4577 neigh=0.4312 | perp c/n=226.8/29.5 | active c/n=256/31 | gini c/n=0.07/0.08 | lr=3.00e-04 gnorm=2.45 | 1.6s 71397 cells/s

[nicheverse] epoch 3/30 | total=0.7604 cell=0.3649 neigh=0.3954 | perp c/n=233.8/29.5 | active c/n=256/31 | gini c/n=0.04/0.07 | lr=3.00e-04 gnorm=2.41 | 1.6s 72235 cells/s

[nicheverse] epoch 4/30 | total=0.7011 cell=0.3269 neigh=0.3741 | perp c/n=234.7/29.5 | active c/n=256/31 | gini c/n=0.05/0.07 | lr=3.00e-04 gnorm=2.37 | 1.6s 71857 cells/s

[nicheverse] epoch 5/30 | total=0.6394 cell=0.3050 neigh=0.3345 | perp c/n=233.5/30.4 | active c/n=256/31 | gini c/n=0.07/0.04 | lr=3.00e-04 gnorm=2.20 | 1.6s 71164 cells/s

[nicheverse] epoch 6/30 | total=0.6143 cell=0.2869 neigh=0.3275 | perp c/n=232.9/30.4 | active c/n=256/31 | gini c/n=0.08/0.04 | lr=3.00e-04 gnorm=2.17 | 1.6s 71745 cells/s

[nicheverse] epoch 7/30 | total=0.5962 cell=0.2764 neigh=0.3199 | perp c/n=231.2/30.5 | active c/n=256/31 | gini c/n=0.10/0.04 | lr=3.00e-04 gnorm=1.99 | 1.6s 71728 cells/s

[nicheverse] epoch 8/30 | total=0.5888 cell=0.2690 neigh=0.3198 | perp c/n=231.2/30.4 | active c/n=256/31 | gini c/n=0.10/0.04 | lr=3.00e-04 gnorm=2.07 | 1.6s 71478 cells/s

[nicheverse] epoch 9/30 | total=0.5792 cell=0.2633 neigh=0.3159 | perp c/n=231.3/30.5 | active c/n=256/31 | gini c/n=0.11/0.04 | lr=3.00e-04 gnorm=1.84 | 1.6s 71846 cells/s

[nicheverse] epoch 10/30 | total=0.5706 cell=0.2597 neigh=0.3109 | perp c/n=232.2/30.6 | active c/n=256/32 | gini c/n=0.10/0.04 | lr=3.00e-04 gnorm=1.70 | 1.6s 71660 cells/s

[nicheverse] epoch 11/30 | total=0.5429 cell=0.2587 neigh=0.2842 | perp c/n=231.6/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.73 | 1.6s 71436 cells/s

[nicheverse] epoch 12/30 | total=0.5422 cell=0.2565 neigh=0.2857 | perp c/n=232.0/31.4 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=2.00 | 1.6s 71549 cells/s

[nicheverse] epoch 13/30 | total=0.5377 cell=0.2545 neigh=0.2831 | perp c/n=232.8/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.80 | 1.6s 72130 cells/s

[nicheverse] epoch 14/30 | total=0.5349 cell=0.2540 neigh=0.2809 | perp c/n=232.8/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.74 | 1.6s 71458 cells/s

[nicheverse] epoch 15/30 | total=0.5322 cell=0.2519 neigh=0.2804 | perp c/n=233.2/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.60 | 1.6s 71309 cells/s

[nicheverse] epoch 16/30 | total=0.5316 cell=0.2520 neigh=0.2796 | perp c/n=232.9/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.61 | 1.6s 71721 cells/s

[nicheverse] epoch 17/30 | total=0.5282 cell=0.2480 neigh=0.2802 | perp c/n=233.7/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.46 | 1.6s 72131 cells/s

[nicheverse] epoch 18/30 | total=0.5274 cell=0.2486 neigh=0.2788 | perp c/n=233.6/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.60 | 1.6s 71383 cells/s

[nicheverse] epoch 19/30 | total=0.5259 cell=0.2478 neigh=0.2781 | perp c/n=234.0/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.70 | 1.6s 71724 cells/s

[nicheverse] epoch 20/30 | total=0.5220 cell=0.2460 neigh=0.2760 | perp c/n=233.5/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.43 | 1.6s 71221 cells/s

[nicheverse] epoch 21/30 | total=0.5184 cell=0.2441 neigh=0.2743 | perp c/n=234.1/31.6 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.35 | 1.6s 71831 cells/s

[nicheverse] epoch 22/30 | total=0.5184 cell=0.2427 neigh=0.2756 | perp c/n=234.7/31.5 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.44 | 1.6s 72288 cells/s

[nicheverse] epoch 23/30 | total=0.5163 cell=0.2430 neigh=0.2733 | perp c/n=234.2/31.6 | active c/n=256/32 | gini c/n=0.09/0.01 | lr=3.00e-04 gnorm=1.27 | 1.6s 72373 cells/s

[nicheverse] epoch 24/30 | total=0.5148 cell=0.2417 neigh=0.2731 | perp c/n=234.0/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.29 | 1.6s 71955 cells/s

[nicheverse] epoch 25/30 | total=0.5137 cell=0.2404 neigh=0.2733 | perp c/n=234.4/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.31 | 1.6s 71889 cells/s

[nicheverse] epoch 26/30 | total=0.5143 cell=0.2417 neigh=0.2725 | perp c/n=233.9/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.32 | 1.6s 72092 cells/s

[nicheverse] epoch 27/30 | total=0.5135 cell=0.2409 neigh=0.2726 | perp c/n=233.6/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.29 | 1.6s 71491 cells/s

[nicheverse] epoch 28/30 | total=0.5145 cell=0.2409 neigh=0.2735 | perp c/n=233.8/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.47 | 1.6s 72352 cells/s

[nicheverse] epoch 29/30 | total=0.5121 cell=0.2397 neigh=0.2724 | perp c/n=234.3/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.32 | 1.6s 71376 cells/s

[nicheverse] epoch 30/30 | total=0.5109 cell=0.2388 neigh=0.2721 | perp c/n=234.3/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.27 | 1.6s 72225 cells/s

What landed in the checkpoint directory#

train_model writes the model, both codebooks, per-cell embeddings and code indices, the loss curve, a runtime record, and an annotated AnnData.

import os, json
print(sorted(os.listdir(CKPT)))
runtime = json.load(open(os.path.join(CKPT, 'training_runtime.json')))
print('\nruntime:', json.dumps(runtime, indent=2))

['adata_with_hierarchical_embeddings.h5ad', 'cell_codebook.npz', 'env_snapshot.json', 'hierarchical_cell_embeddings.npz', 'hierarchical_cell_indices.npz', 'hierarchical_neighborhood_embeddings.npz', 'hierarchical_neighborhood_indices.npz', 'hierarchical_vqvae_checkpoint.json', 'hierarchical_vqvae_checkpoint.pt', 'neighborhood_codebook.npz', 'train_config.json', 'training_curves.pdf', 'training_losses.json', 'training_metrics.csv', 'training_runtime.json']

runtime: {
  "total_seconds": 48.046,
  "total_hms": "0:00:48",
  "n_epochs": 30,
  "mean_epoch_seconds": 1.602,
  "cells_per_second": 70797.68,
  "iters_per_second": 34.966,
  "n_cells": 113385,
  "effective_batch_size": 2048,
  "nominal_learning_rate": 0.0003,
  "effective_learning_rate": 0.0003,
  "lr_scaled_with_batch": false,
  "n_batches_per_epoch": 56,
  "peak_gpu_gb": 1.154,
  "device": "cuda",
  "encoder_type": "mlp_deep",
  "quantizer_type": "vq",
  "input_dim": 368
}

The codes are attached back onto the AnnData: obs['cell_codebook_idx'] (0..255), obs['neighborhood_codebook_idx'] (0..31), and the continuous embeddings in obsm.

print(nv.anndata_keys())
cell_idx = adata.obs['cell_codebook_idx'].to_numpy()
neigh_idx = adata.obs['neighborhood_codebook_idx'].to_numpy()
print('cell codes used:', len(np.unique(cell_idx)), '/', mc.cell_num_embeddings)
print('niches used:', len(np.unique(neigh_idx)), '/', mc.neighborhood_num_embeddings)
print('embedding shapes:', adata.obsm['X_cell_embedding'].shape, adata.obsm['X_neighborhood_embedding'].shape)

{'cell_code': 'cell_codebook_idx', 'neighborhood_code': 'neighborhood_codebook_idx', 'sample': 'sample_id', 'cell_embedding': 'X_cell_embedding', 'neighborhood_embedding': 'X_neighborhood_embedding', 'spatial': 'spatial'}
cell codes used: 256 / 256
niches used: 32 / 32
embedding shapes: (113385, 64) (113385, 256)

Per-code top markers#

To read the cell codebook, take for each code the genes most enriched in cells assigned to it (mean log-normalized expression, z-scored across codes). This is a first-pass summary; the package provides a full annotation workflow (nicheverse.annotate) with DEGs, site distribution, and literature grounding.

# log-normalize a copy for the marker summary (train_model normalized internally,
# but the returned adata.X here is already log-normalized)
expr = adata.copy()
if 'log1p' not in expr.uns:
    sc.pp.normalize_total(expr); sc.pp.log1p(expr)
X = expr.X.toarray() if hasattr(expr.X, 'toarray') else np.asarray(expr.X)
genes = np.array(expr.var_names.astype(str))
codes = np.unique(cell_idx)
mean_by_code = np.vstack([X[cell_idx == k].mean(0) for k in codes])
z = (mean_by_code - mean_by_code.mean(0)) / (mean_by_code.std(0) + 1e-8)
rows = []
for i, k in enumerate(codes):
    top = genes[np.argsort(z[i])[::-1][:5]]
    rows.append({'code': int(k), 'n_cells': int((cell_idx == k).sum()), 'top_markers': ', '.join(top)})
marker_table = pd.DataFrame(rows).sort_values('n_cells', ascending=False)
marker_table.head(15)
code n_cells top_markers
85 85 774 Nrl, Lima1, C1ql3, Wnt5a, Cdhr1
188 188 752 Hapln1, Nrl, 4833423E24Rik, Tagln2, Rax
52 52 698 Wnt5a, Calca, Slc17a7, Nr2e3, Camta1
21 21 683 Pde6c, Mylk, Gulo, Cngb3, Kcne2
177 177 665 Gpr83, Cdhr1, Lima1, Wnt5a, Reep6
133 133 657 Igfn1, Vsx1, Grm6, Lect1, Reln
155 155 655 Pappa2, Slc18a3, Col25a1, Gng7, Mmp17
183 183 643 Etv1, Igfbp4, Fgf7, Sox4, Camta1
31 31 637 Antxr2, Rnf152, Kcne2, Prkca, Tpbg
164 164 629 Tax1bp1, Penk, Dmrtb1, Gpr83, Nt5e
180 180 627 Mfap5, Crybb3, 4833423E24Rik, Rarres1, Slc17a7
163 163 623 Lect1, Reln, Vsx1, Il1rapl2, Igfn1
201 201 622 Igfn1, Il1rapl2, Vsx1, St18, Grm6
73 73 618 Prokr1, Mylk, Pde6c, Gulo, Tsc22d4
225 225 617 St18, Ngfr, Isl1, Ddit4l, Nfia

Code-usage bar chart#

How evenly the cells spread across the cell codebook. A well-utilized codebook occupies most codes; occupancy grows with the biological diversity and scale of the dataset.

counts = pd.Series(cell_idx).value_counts().sort_values(ascending=False)
fig, ax = plt.subplots(figsize=(6, 3))
ax.bar(range(len(counts)), counts.values, width=1.0)
ax.set_xlabel('cell code (sorted by usage)'); ax.set_ylabel('n cells'); ax.set_yscale('log')
ax.set_title(f'cell codebook usage ({len(counts)}/{mc.cell_num_embeddings} codes used)')
fig.tight_layout(); plt.show()
print('done')

done

Next steps#

  • 02_transcript_context.ipynb adds the segmentation-free molecular field.

  • 03_molecule_set.ipynb uses the subcellular transcript point cloud.

  • 04_apply_to_new_data.ipynb assigns this codebook to a held-out sample.

For real cohorts set num_epochs ~300 and annotate the codebook with nicheverse.annotate (per-code DEGs + literature grounding).

On this page