02 Transcript context: the segmentation-free molecular field#

Segmentation assigns each transcript to one cell. Transcript context instead counts every molecule within a small radius of a cell’s centroid, regardless of segmentation, giving a local molecular field that is robust to segmentation errors and captures the immediate microenvironment. We concatenate it onto the segmented counts, doubling the input dimension, and train the same hierarchical VQ-VAE.

Here we use the bundled real Xenium RCC TMA core (7,824 cells, 366-gene panel), which ships with its matched molecule table.

Demo note: ~20 epochs for speed; a real cohort uses num_epochs ~300. The default radius is 7 um (measured to capture roughly one cell from its nucleus center with minimal neighbor bleed).

import numpy as np, pandas as pd, scanpy as sc, anndata as ad
import nicheverse as nv
from nicheverse import ModelConfig, TrainConfig
from nicheverse.data import transcript_context
import os
DATA = os.path.join('..', 'examples', 'data')
adata = nv.read_spatial(f'{DATA}/xenium_rcc_core.h5ad', sample_col='sample_id')
print(adata.n_obs, 'cells x', adata.n_vars, 'genes')

7824 cells x 366 genes

Compute the transcript-context field#

transcript_context reads a per-sample molecule table (x_location, y_location, feature_name for Xenium), drops control / blank probes, and for each cell counts the molecules within radius microns of its centroid, returning a log1p (n_cells, n_genes) matrix. Pass a single path for a single-sample run, or a {sample_id: path} mapping for a cohort.

tx_path = f'{DATA}/xenium_rcc_core_transcripts.parquet'
txc = transcript_context(adata, tx_path, radius=7.0, sample_col='sample_id', platform='xenium')
print('transcript-context matrix:', txc.shape)
print('mean non-zero genes per cell:', float((txc > 0).sum(1).mean()).__round__(1))

transcript-context matrix: (7824, 366)
mean non-zero genes per cell: 67.7

Concatenate onto counts and train#

The model input is concat(log-norm counts, transcript context), so input_dim doubles to 732. We log-normalize the counts ourselves and disable normalize / log1p in TrainConfig (the context field is already log1p). Because the input is already log-normalized rather than raw counts, we also select the MSE reconstruction (cell_recon='mse', niche_recon='mse', detection_weight=0) instead of the count-likelihood default, which expects raw integer counts. We keep gene_names=() because the concatenated matrix is not a plain gene panel. A good encoder for transcript context is mlp_deep (the library default).

expr = adata.copy(); sc.pp.normalize_total(expr); sc.pp.log1p(expr)
X_counts = expr.X.toarray() if hasattr(expr.X, 'toarray') else np.asarray(expr.X)
X_in = np.concatenate([X_counts, txc], axis=1).astype(np.float32)
train_ad = ad.AnnData(X=X_in, obs=adata.obs.copy(),
                      obsm={'spatial': np.asarray(adata.obsm['spatial'])})
train_ad.obs_names = adata.obs_names
train_ad.uns['log1p'] = {'base': None}   # mark as already log-normalized

mc = ModelConfig(input_dim=X_in.shape[1], gene_names=(), encoder_type='mlp_deep',
                 cell_recon='mse', niche_recon='mse', detection_weight=0,
                 cell_num_embeddings=256, neighborhood_num_embeddings=32)
# reference defaults for graph/seed; cell_recon/niche_recon set to 'mse' because the
# transcript-context input is already log-normalized (not raw counts). num_epochs
# (300 -> 20) and batch_size (32768 -> 2048) are shrunk for the demo.
tc = TrainConfig(num_epochs=20, batch_size=2048, k_neighbors=20,
                 spatial_graph='knn_radius', radius=50.0,
                 normalize=False, log1p=False, save_best=False, seed=9)
model, out = nv.train_model(train_ad, 'runs/transcript_context',
                            model_config=mc, train_config=tc, sample_col='sample_id')
print('input_dim (counts + context):', X_in.shape[1])

[nicheverse] epoch 1/20 | total=2.2960 cell=1.1563 neigh=1.1397 | perp c/n=115.7/15.4 | active c/n=256/32 | gini c/n=0.57/0.48 | lr=3.00e-04 gnorm=13.81 | 0.5s 14411 cells/s

[nicheverse] epoch 2/20 | total=2.2098 cell=1.2059 neigh=1.0039 | perp c/n=105.7/16.5 | active c/n=256/31 | gini c/n=0.64/0.52 | lr=3.00e-04 gnorm=12.37 | 0.1s 69022 cells/s

[nicheverse] epoch 3/20 | total=1.9165 cell=1.0468 neigh=0.8697 | perp c/n=121.7/18.3 | active c/n=256/32 | gini c/n=0.59/0.49 | lr=3.00e-04 gnorm=10.58 | 0.1s 70546 cells/s

[nicheverse] epoch 4/20 | total=1.7096 cell=0.9112 neigh=0.7985 | perp c/n=136.0/19.0 | active c/n=251/31 | gini c/n=0.55/0.49 | lr=3.00e-04 gnorm=8.56 | 0.1s 72533 cells/s

[nicheverse] epoch 5/20 | total=1.5407 cell=0.8363 neigh=0.7043 | perp c/n=146.0/20.4 | active c/n=254/31 | gini c/n=0.53/0.46 | lr=3.00e-04 gnorm=7.88 | 0.1s 70001 cells/s

[nicheverse] epoch 6/20 | total=1.3014 cell=0.7055 neigh=0.5959 | perp c/n=166.9/22.6 | active c/n=256/32 | gini c/n=0.46/0.37 | lr=3.00e-04 gnorm=7.05 | 0.1s 72749 cells/s

[nicheverse] epoch 7/20 | total=1.1987 cell=0.6151 neigh=0.5837 | perp c/n=179.6/22.5 | active c/n=256/32 | gini c/n=0.41/0.41 | lr=3.00e-04 gnorm=6.43 | 0.1s 72847 cells/s

[nicheverse] epoch 8/20 | total=1.0764 cell=0.5425 neigh=0.5339 | perp c/n=191.5/23.2 | active c/n=256/32 | gini c/n=0.36/0.39 | lr=3.00e-04 gnorm=4.73 | 0.1s 67901 cells/s

[nicheverse] epoch 9/20 | total=0.9581 cell=0.4866 neigh=0.4715 | perp c/n=203.0/24.5 | active c/n=256/32 | gini c/n=0.33/0.32 | lr=3.00e-04 gnorm=3.90 | 0.1s 72556 cells/s

[nicheverse] epoch 10/20 | total=0.8741 cell=0.4463 neigh=0.4278 | perp c/n=207.8/25.3 | active c/n=256/32 | gini c/n=0.29/0.28 | lr=3.00e-04 gnorm=3.89 | 0.1s 72507 cells/s

[nicheverse] epoch 11/20 | total=0.7809 cell=0.4035 neigh=0.3774 | perp c/n=216.3/26.4 | active c/n=256/32 | gini c/n=0.25/0.27 | lr=3.00e-04 gnorm=3.22 | 0.1s 72721 cells/s

[nicheverse] epoch 12/20 | total=0.7173 cell=0.3749 neigh=0.3424 | perp c/n=219.6/27.2 | active c/n=256/32 | gini c/n=0.23/0.23 | lr=3.00e-04 gnorm=3.16 | 0.1s 70884 cells/s

[nicheverse] epoch 13/20 | total=0.6502 cell=0.3456 neigh=0.3046 | perp c/n=224.9/28.0 | active c/n=256/32 | gini c/n=0.20/0.21 | lr=3.00e-04 gnorm=2.90 | 0.2s 40057 cells/s

[nicheverse] epoch 14/20 | total=0.5822 cell=0.3199 neigh=0.2623 | perp c/n=229.6/29.0 | active c/n=256/32 | gini c/n=0.16/0.17 | lr=3.00e-04 gnorm=2.84 | 0.1s 67819 cells/s

[nicheverse] epoch 15/20 | total=0.5421 cell=0.3082 neigh=0.2340 | perp c/n=230.7/29.7 | active c/n=256/32 | gini c/n=0.16/0.14 | lr=3.00e-04 gnorm=2.18 | 0.1s 69396 cells/s

[nicheverse] epoch 16/20 | total=0.5054 cell=0.2930 neigh=0.2125 | perp c/n=233.1/30.3 | active c/n=256/32 | gini c/n=0.14/0.12 | lr=3.00e-04 gnorm=1.94 | 0.1s 72683 cells/s

[nicheverse] epoch 17/20 | total=0.4833 cell=0.2852 neigh=0.1981 | perp c/n=231.7/30.5 | active c/n=256/31 | gini c/n=0.14/0.09 | lr=3.00e-04 gnorm=1.81 | 0.1s 71240 cells/s

[nicheverse] epoch 18/20 | total=0.4677 cell=0.2735 neigh=0.1942 | perp c/n=234.6/30.6 | active c/n=256/32 | gini c/n=0.13/0.10 | lr=3.00e-04 gnorm=1.74 | 0.1s 71698 cells/s

[nicheverse] epoch 19/20 | total=0.4532 cell=0.2636 neigh=0.1896 | perp c/n=235.1/30.6 | active c/n=256/32 | gini c/n=0.12/0.09 | lr=3.00e-04 gnorm=1.58 | 0.1s 72576 cells/s

[nicheverse] epoch 20/20 | total=0.4406 cell=0.2583 neigh=0.1823 | perp c/n=235.3/30.6 | active c/n=256/32 | gini c/n=0.11/0.09 | lr=3.00e-04 gnorm=1.62 | 0.1s 69432 cells/s

input_dim (counts + context): 732

cell_idx = out.obs['cell_codebook_idx'].to_numpy()
print('cell codes used:', len(np.unique(cell_idx)), '/ 256')
print('niches used:', out.obs['neighborhood_codebook_idx'].nunique(), '/ 32')
print('embeddings:', out.obsm['X_cell_embedding'].shape)

cell codes used: 256 / 256
niches used: 32 / 32
embeddings: (7824, 64)

This single homogeneous TMA core exercises only a few cell codes; codebook fullness grows with the diversity and scale of a real multi-sample cohort (see notebook 01 on MERFISH). The point here is the transcript-context representation and how it is concatenated and trained. To store both, keep .X as the 366-gene expression and stash the 732-dim model input in obsm for reproducibility.

On this page