03 Molecule set: the subcellular transcript point cloud#

The richest input keeps every transcript as a point: for each cell, the set of molecules within a radius of its centroid, each carrying its gene identity and a subcellular offset (dx, dy) in microns. A permutation-invariant Set-Transformer encoder (MoleculeSetVQVAE) pools this point cloud into a cell embedding, quantized against the same cell codebook, with an aggregated kNN neighborhood branch fused by cross-attention.

We use the bundled real Xenium RCC core; its molecule-set shards were built offline from the transcripts table (radius 7 um).

Demo note: ~8 epochs for speed; real runs use many more. The Fourier positional encoding and gather radius both use max_micron = 7.0.

import numpy as np, scanpy as sc
import torch, torch.nn.functional as F
from torch.utils.data import DataLoader
import nicheverse as nv
from nicheverse.data import MoleculeSetDataset
from nicheverse.models.molecule_set import MoleculeSetVQVAE
import os
DATA = os.path.join('..', 'examples', 'data')
adata = nv.read_spatial(f'{DATA}/xenium_rcc_core.h5ad', sample_col='sample_id')
G = adata.n_vars
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('cells', adata.n_obs, 'genes', G, 'device', device)

cells 7824 genes 366 device cuda

Build the molecule-set dataset#

MoleculeSetDataset aligns the per-cell molecule-set shards to the AnnData’s obs_names and, with with_neighborhood=True, attaches the aggregated kNN neighborhood feature (it needs the log-normalized expression, the spatial coordinates, and the sample labels). Each item is (gene, coords, mask, comp, neigh, idx) where comp is the log1p within-radius gene composition used as the cell reconstruction target.

expr = adata.copy(); sc.pp.normalize_total(expr); sc.pp.log1p(expr)
X_lognorm = expr.X.toarray() if hasattr(expr.X, 'toarray') else np.asarray(expr.X)
X_lognorm = X_lognorm.astype(np.float32)
ds = MoleculeSetDataset(
    adata.obs_names.astype(str),
    shard_dir=f'{DATA}/xenium_rcc_core_molecule_sets',
    spatial_coords=np.asarray(adata.obsm['spatial']),
    sample_ids=adata.obs['sample_id'].astype(str).to_numpy(),
    expr_lognorm=X_lognorm,
    k_neighbors=20, with_neighborhood=True,
)
print('molecules per cell (padded length M):', ds.M, ' genes G:', ds.G, ' cells:', len(ds))

molecules per cell (padded length M): 128  genes G: 366  cells: 7824

Train the molecule-set VQ-VAE#

MoleculeSetVQVAE emits the same cell / neighborhood code indices as the standard model, so downstream annotation is unchanged. We train a short loop by hand to show the interface.

torch.manual_seed(9)  # default seed
model = MoleculeSetVQVAE(n_genes=G, cell_num_embeddings=256, neighborhood_num_embeddings=32,
                         use_neighborhood=True, use_cross_attention=True).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
loader = DataLoader(ds, batch_size=1024, shuffle=True)
for ep in range(8):
    model.train(); tot = 0.0; nb = 0
    for gene, coords, mask, comp, neigh, _ in loader:
        if gene.shape[0] < 2:
            continue
        gene, coords, mask = gene.to(device), coords.to(device), mask.to(device)
        comp, neigh = comp.to(device), neigh.to(device)
        opt.zero_grad(set_to_none=True)
        cr, nr, cvq, nvq, ci, ni, cp, npp = model(gene, coords, mask, neigh)
        loss = F.mse_loss(cr, comp) + cvq + F.mse_loss(nr, neigh) + nvq
        loss.backward(); opt.step()
        tot += float(loss); nb += 1
    print(f'epoch {ep+1}/8 loss={tot/max(nb,1):.4f} cell_perplexity={float(cp):.1f}')

epoch 1/8 loss=1.9090 cell_perplexity=2.4

epoch 2/8 loss=1.4716 cell_perplexity=2.0

epoch 3/8 loss=1.2657 cell_perplexity=2.8

epoch 4/8 loss=1.0618 cell_perplexity=2.9

epoch 5/8 loss=0.9414 cell_perplexity=3.5

epoch 6/8 loss=0.8637 cell_perplexity=3.8

epoch 7/8 loss=0.8246 cell_perplexity=46.6

epoch 8/8 loss=0.7129 cell_perplexity=26.5

Assign codes to every cell#

model.eval(); cidx = []
with torch.inference_mode():
    for gene, coords, mask, comp, neigh, _ in DataLoader(ds, batch_size=1024, shuffle=False):
        gene, coords, mask, neigh = gene.to(device), coords.to(device), mask.to(device), neigh.to(device)
        _, _, _, _, ci, ni, _, _ = model(gene, coords, mask, neigh)
        cidx.append(ci.cpu().numpy().reshape(-1))
cidx = np.concatenate(cidx)
print('cells coded:', len(cidx), ' cell codes used:', len(np.unique(cidx)), '/ 256')
adata.obs['cell_codebook_idx'] = cidx
print('done')

cells coded: 7824  cell codes used: 80 / 256
done

The subcellular point cloud is the most information-rich input (it keeps molecular position within the cell). On a diverse cohort it produces a healthy, biologically coherent codebook; the fullness on this single homogeneous core is limited by biology, not the method. Important: for any set / attention encoder, pool with concat[masked max, masked mean, PMA] (PMA alone collapses the pooled embedding); MoleculeSetEncoder already does this.

On this page