04 Apply a trained model to new data#
Once a codebook is trained, predict_codes assigns the same cell and neighborhood codes to a new sample without retraining. This is how you annotate a held-out slide with an existing codebook.
We train on three of the four MERFISH retina samples and apply the codebook to the held-out fourth. The neighbor-graph arguments passed to predict_codes must match the ones used at training time.
import numpy as np
import nicheverse as nv
from nicheverse import ModelConfig, TrainConfig, predict_codes
import os
MERFISH = os.path.join('..', 'examples', 'data', 'merfish_retina.h5ad')
adata = nv.read_spatial(MERFISH, sample_col='sample_id')
samples = sorted(adata.obs['sample_id'].astype(str).unique())
held_out = samples[-1]
train_samples = samples[:-1]
print('train on:', train_samples, ' held out:', held_out)
train on: ['VZG105a_WT1', 'VZG105a_WT2', 'VZG105a_WT3'] held out: VZG105a_WT4
is_train = adata.obs['sample_id'].astype(str).isin(train_samples)
train_ad = adata[is_train].copy()
new_ad = adata[~is_train].copy()
print('train cells:', train_ad.n_obs, ' held-out cells:', new_ad.n_obs)
train cells: 86675 held-out cells: 26710
Train the codebook#
The graph settings are captured once and reused for both train and predict so the neighborhood codes stay comparable.
# reference defaults (encoder mlp_deep, vq, knn_radius@50um, k=20, weighted_mean, seed 9);
# only num_epochs (300 -> 30) and batch_size (32768 -> 2048) are shrunk for the demo.
GRAPH = dict(k_neighbors=20, neighborhood_aggregation='weighted_mean',
spatial_graph='knn_radius', radius=50.0)
mc = ModelConfig(input_dim=train_ad.n_vars, gene_names=tuple(train_ad.var_names.astype(str)),
encoder_type='mlp_deep', cell_num_embeddings=256, neighborhood_num_embeddings=32)
tc = TrainConfig(num_epochs=30, batch_size=2048, save_best=False, seed=9, **GRAPH)
CKPT = 'runs/apply_demo'
model, train_ad = nv.train_model(train_ad, CKPT, model_config=mc, train_config=tc, sample_col='sample_id')
print('trained; codes used:', train_ad.obs['cell_codebook_idx'].nunique(), '/ 256')
[nicheverse] epoch 1/30 | total=1.7807 cell=0.9873 neigh=0.7933 | perp c/n=166.9/23.8 | active c/n=256/32 | gini c/n=0.24/0.18 | lr=3.00e-04 gnorm=6.79 | 1.6s 53312 cells/s
[nicheverse] epoch 2/30 | total=1.0496 cell=0.5214 neigh=0.5282 | perp c/n=222.1/27.6 | active c/n=253/29 | gini c/n=0.10/0.13 | lr=3.00e-04 gnorm=2.79 | 1.3s 66846 cells/s
[nicheverse] epoch 3/30 | total=0.8895 cell=0.4092 neigh=0.4804 | perp c/n=232.4/28.1 | active c/n=254/29 | gini c/n=0.05/0.12 | lr=3.00e-04 gnorm=2.53 | 1.2s 71212 cells/s
[nicheverse] epoch 4/30 | total=0.8060 cell=0.3637 neigh=0.4422 | perp c/n=234.1/28.5 | active c/n=255/29 | gini c/n=0.05/0.10 | lr=3.00e-04 gnorm=2.14 | 1.2s 71181 cells/s
[nicheverse] epoch 5/30 | total=0.7674 cell=0.3376 neigh=0.4298 | perp c/n=233.5/28.5 | active c/n=255/29 | gini c/n=0.06/0.10 | lr=3.00e-04 gnorm=2.38 | 1.2s 71051 cells/s
[nicheverse] epoch 6/30 | total=0.7404 cell=0.3207 neigh=0.4197 | perp c/n=232.6/28.5 | active c/n=255/29 | gini c/n=0.08/0.10 | lr=3.00e-04 gnorm=2.14 | 1.2s 71283 cells/s
[nicheverse] epoch 7/30 | total=0.7193 cell=0.3086 neigh=0.4106 | perp c/n=230.8/28.6 | active c/n=255/29 | gini c/n=0.09/0.10 | lr=3.00e-04 gnorm=1.92 | 1.2s 71269 cells/s
[nicheverse] epoch 8/30 | total=0.7098 cell=0.3011 neigh=0.4087 | perp c/n=229.4/28.5 | active c/n=255/29 | gini c/n=0.11/0.10 | lr=3.00e-04 gnorm=2.14 | 1.2s 71335 cells/s
[nicheverse] epoch 9/30 | total=0.6960 cell=0.2900 neigh=0.4061 | perp c/n=229.9/28.5 | active c/n=255/29 | gini c/n=0.12/0.10 | lr=3.00e-04 gnorm=1.84 | 1.2s 71789 cells/s
[nicheverse] epoch 10/30 | total=0.6861 cell=0.2849 neigh=0.4012 | perp c/n=228.8/28.6 | active c/n=255/29 | gini c/n=0.13/0.10 | lr=3.00e-04 gnorm=1.78 | 1.2s 71680 cells/s
[nicheverse] epoch 11/30 | total=0.6841 cell=0.2827 neigh=0.4014 | perp c/n=227.4/28.5 | active c/n=255/29 | gini c/n=0.13/0.10 | lr=3.00e-04 gnorm=1.86 | 1.2s 70417 cells/s
[nicheverse] epoch 12/30 | total=0.6810 cell=0.2801 neigh=0.4009 | perp c/n=227.9/28.5 | active c/n=255/29 | gini c/n=0.13/0.10 | lr=3.00e-04 gnorm=1.96 | 1.2s 71432 cells/s
[nicheverse] epoch 13/30 | total=0.6647 cell=0.2751 neigh=0.3896 | perp c/n=228.9/28.9 | active c/n=256/32 | gini c/n=0.13/0.09 | lr=3.00e-04 gnorm=2.04 | 1.2s 71057 cells/s
[nicheverse] epoch 14/30 | total=0.5805 cell=0.2733 neigh=0.3072 | perp c/n=228.6/31.4 | active c/n=256/32 | gini c/n=0.13/0.01 | lr=3.00e-04 gnorm=1.91 | 1.2s 71633 cells/s
[nicheverse] epoch 15/30 | total=0.5651 cell=0.2672 neigh=0.2978 | perp c/n=230.5/31.5 | active c/n=256/32 | gini c/n=0.12/0.01 | lr=3.00e-04 gnorm=1.43 | 1.2s 71504 cells/s
[nicheverse] epoch 16/30 | total=0.5663 cell=0.2672 neigh=0.2991 | perp c/n=230.9/31.5 | active c/n=256/32 | gini c/n=0.12/0.01 | lr=3.00e-04 gnorm=1.66 | 1.2s 71806 cells/s
[nicheverse] epoch 17/30 | total=0.5603 cell=0.2645 neigh=0.2958 | perp c/n=232.0/31.5 | active c/n=256/32 | gini c/n=0.12/0.01 | lr=3.00e-04 gnorm=1.41 | 1.2s 71659 cells/s
[nicheverse] epoch 18/30 | total=0.5601 cell=0.2627 neigh=0.2975 | perp c/n=232.1/31.5 | active c/n=256/32 | gini c/n=0.11/0.01 | lr=3.00e-04 gnorm=1.56 | 1.2s 71269 cells/s
[nicheverse] epoch 19/30 | total=0.5562 cell=0.2625 neigh=0.2936 | perp c/n=232.5/31.6 | active c/n=256/32 | gini c/n=0.11/0.01 | lr=3.00e-04 gnorm=1.33 | 1.2s 71620 cells/s
[nicheverse] epoch 20/30 | total=0.5565 cell=0.2628 neigh=0.2937 | perp c/n=232.5/31.6 | active c/n=256/32 | gini c/n=0.11/0.01 | lr=3.00e-04 gnorm=1.48 | 1.2s 72732 cells/s
[nicheverse] epoch 21/30 | total=0.5571 cell=0.2608 neigh=0.2963 | perp c/n=232.8/31.4 | active c/n=256/32 | gini c/n=0.11/0.01 | lr=3.00e-04 gnorm=1.66 | 1.2s 71350 cells/s
[nicheverse] epoch 22/30 | total=0.5546 cell=0.2607 neigh=0.2939 | perp c/n=233.0/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.50 | 1.2s 71714 cells/s
[nicheverse] epoch 23/30 | total=0.5525 cell=0.2590 neigh=0.2935 | perp c/n=232.9/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.42 | 1.2s 72168 cells/s
[nicheverse] epoch 24/30 | total=0.5511 cell=0.2588 neigh=0.2923 | perp c/n=233.2/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.41 | 1.2s 71678 cells/s
[nicheverse] epoch 25/30 | total=0.5495 cell=0.2573 neigh=0.2922 | perp c/n=233.3/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.45 | 1.2s 71560 cells/s
[nicheverse] epoch 26/30 | total=0.5501 cell=0.2569 neigh=0.2932 | perp c/n=233.0/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.50 | 1.2s 71490 cells/s
[nicheverse] epoch 27/30 | total=0.5471 cell=0.2567 neigh=0.2903 | perp c/n=233.2/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.28 | 1.2s 71627 cells/s
[nicheverse] epoch 28/30 | total=0.5463 cell=0.2557 neigh=0.2906 | perp c/n=233.3/31.6 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.23 | 1.2s 71743 cells/s
[nicheverse] epoch 29/30 | total=0.5468 cell=0.2559 neigh=0.2909 | perp c/n=233.1/31.5 | active c/n=256/32 | gini c/n=0.10/0.01 | lr=3.00e-04 gnorm=1.38 | 1.2s 71013 cells/s
[nicheverse] epoch 30/30 | total=0.5449 cell=0.2554 neigh=0.2895 | perp c/n=232.8/31.6 | active c/n=256/32 | gini c/n=0.11/0.01 | lr=3.00e-04 gnorm=1.36 | 1.2s 71242 cells/s
trained; codes used: 256 / 256
Assign codes to the held-out sample#
predict_codes loads the checkpoint, aligns the gene panel, builds the neighbor graph the same way, and writes cell_codebook_idx / neighborhood_codebook_idx plus the embeddings onto the new AnnData. It never updates the codebook.
ckpt_pt = f'{CKPT}/hierarchical_vqvae_checkpoint.pt'
coded = predict_codes(new_ad, ckpt_pt, sample_col='sample_id', **GRAPH,
output_path='runs/held_out_annotated.h5ad')
c = coded.obs['cell_codebook_idx'].to_numpy()
print('held-out cells coded:', coded.n_obs)
print('distinct cell codes on held-out sample:', len(np.unique(c)))
print('niches on held-out sample:', coded.obs['neighborhood_codebook_idx'].nunique())
assert 'X_cell_embedding' in coded.obsm
print('embeddings attached:', coded.obsm['X_cell_embedding'].shape)
held-out cells coded: 26710
distinct cell codes on held-out sample: 256
niches on held-out sample: 32
embeddings attached: (26710, 64)
The held-out sample now carries the same codes as the training cohort, so a cell assigned code 42 in training and code 42 here refer to the same learned state. That is what makes the codebook a shared, transferable annotation across samples and studies.