05 Annotate codes into cell types and spatial niches#

Training gives every cell an integer cell code and niche code. This notebook turns those anonymous codes into named cell types and named niches with the nicheverse.annotate workflow.

The pipeline has two parts:

  1. Evidence (offline, no API key). For each code, nicheverse.annotate assembles per-code markers (z-scored mean expression), one-vs-rest differential expression, the code’s distribution over any metadata you pass, and a coarse-to-fine grouping of the codes. This is deterministic and runs anywhere.

  2. Labeling (optional, needs an LLM provider). An LLM reads that evidence, grounds each label in the markers, DEGs, and primary literature, and returns a ranked label with a confidence and citations; a verification pass, an adversarial pass, and a cross-code reconciliation pass follow, and low-confidence codes are routed to expert review.

Steps 1 runs out of the box; step 2 runs only if a provider key is set, so this notebook is fully runnable offline (it prints the labeling call and skips it).

import os
import numpy as np, pandas as pd
import nicheverse as nv
from nicheverse import ModelConfig, TrainConfig
from nicheverse.annotate import code_evidence, cluster_codes, annotate_codes, attach_labels, annotate_niches
DATA = os.path.join('..', 'examples', 'data')
print('nicheverse', nv.__version__)

nicheverse 0.2.0

Train a quick codebook#

We reuse the quickstart recipe on the bundled MERFISH retina cohort (a diverse, gene-rich dataset that exercises many cell codes). Demo epochs for speed; a real cohort uses num_epochs ~300.

adata = nv.read_spatial(os.path.join(DATA, 'merfish_retina.h5ad'), sample_col='sample_id')
mc = ModelConfig(input_dim=adata.n_vars, gene_names=tuple(adata.var_names.astype(str)),
                 encoder_type='mlp_deep', cell_num_embeddings=256, neighborhood_num_embeddings=32)
tc = TrainConfig(num_epochs=15, batch_size=2048, spatial_graph='knn_radius', radius=50.0,
                 k_neighbors=20, save_best=False, seed=9)
model, adata = nv.train_model(adata, 'runs/annotate_demo', model_config=mc, train_config=tc,
                              sample_col='sample_id')
print('cell codes used:', adata.obs['cell_codebook_idx'].nunique(), '/ 256')
print('niches used:', adata.obs['neighborhood_codebook_idx'].nunique(), '/ 32')

[nicheverse] epoch 1/15 | 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.1s 54447 cells/s

[nicheverse] epoch 2/15 | 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 71811 cells/s

[nicheverse] epoch 3/15 | 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 70148 cells/s

[nicheverse] epoch 4/15 | 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 71565 cells/s

[nicheverse] epoch 5/15 | 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 71439 cells/s

[nicheverse] epoch 6/15 | 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 71924 cells/s

[nicheverse] epoch 7/15 | 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 70909 cells/s

[nicheverse] epoch 8/15 | 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 71215 cells/s

[nicheverse] epoch 9/15 | 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 71474 cells/s

[nicheverse] epoch 10/15 | 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 71442 cells/s

[nicheverse] epoch 11/15 | 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 71718 cells/s

[nicheverse] epoch 12/15 | 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 70768 cells/s

[nicheverse] epoch 13/15 | 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 71010 cells/s

[nicheverse] epoch 14/15 | 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 71479 cells/s

[nicheverse] epoch 15/15 | 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 71326 cells/s

cell codes used: 256 / 256
niches used: 32 / 32

Step 1a: per-code evidence (offline)#

code_evidence returns, for every cell code, its cell count, its top markers by z-score across codes, its one-vs-rest DEGs, and its distribution over any obs columns you name. This is exactly the evidence the LLM (or a human) reads to name the code.

ev = code_evidence(adata, 'cell_codebook_idx', extra_cols=('sample_id',), top_markers=8, top_degs=8)
rows = []
for code, e in ev.items():
    markers = ', '.join(g for g, _z in e['top_markers'][:6])
    rows.append({'cell_code': code, 'n_cells': e['n_cells'],
                 'frac': round(e['frac'], 3), 'top_markers': markers})
evidence_tbl = pd.DataFrame(rows).sort_values('n_cells', ascending=False).reset_index(drop=True)
evidence_tbl.head(12)
cell_code n_cells frac top_markers
0 174 792 0.007 Prkca, Casp7, Tpbg, Strip2, Vstm2b, Rnf152
1 83 748 0.007 Fezf2, Cpne7, Grik1, Otor, Slc16a7, Glra1
2 71 744 0.007 4833423E24Rik, Casz1, Cdhr1, Crx, Purb, Gas7
3 154 686 0.006 Prkca, Casp7, Vstm2b, Adrb1, Tpbg, Slc7a11
4 198 686 0.006 Thrb, Gas7, Casz1, Lyar, Pde6a, Crx
5 222 676 0.006 Meis2, Nos1, Cacna2d2, Gm12371, Cbln2, Tfap2e
6 63 676 0.006 Lmo3, Pde6a, Casz1, Ramp3, Cdhr1, Ybx3
7 136 670 0.006 Drd4, 4833423E24Rik, Inadl, Pde6a, Reep6, Slc17a7
8 124 667 0.006 Drd4, Cngb1, Inadl, Nt5e, Reep6, Prom1
9 29 667 0.006 Nt5e, Pde6a, Ybx3, Ramp3, Gas7, Casz1
10 246 657 0.006 Ngf, Pcsk9, Lhx9, Cbln2, Igfbp2, Olfm3
11 9 656 0.006 Fezf2, Otor, Cpne7, Lrrtm3, Grik1, Slc16a7

Step 1b: coarse-to-fine grouping (offline)#

cluster_codes groups codes by the correlation of their expression profiles, so you can review a hierarchy (broad lineage -> fine state) instead of 256 codes at once. This mirrors how the manuscript pipeline reviews the codebook.

groups = cluster_codes(adata, 'cell_codebook_idx', n_clusters=8)
print(groups.columns.tolist())
groups.head(12)

['cluster']
cluster
0 1
1 1
2 5
3 1
4 5
5 7
6 7
7 1
8 5
9 5
10 1
11 1

Step 2: label the codes with an LLM (optional)#

annotate_codes sends each code’s evidence to a provider (anthropic, openai, or a local ollama) and returns a table with a label, compartment, confidence, rationale, key_markers, citations, and a cross-code reconciled label_refined. It needs the matching key (ANTHROPIC_API_KEY / OPENAI_API_KEY) and the [llm] extra (pip install "nicheverse[llm]").

The cell below runs it only if a key is present; otherwise it prints the call so the notebook stays runnable offline.

have_key = bool(os.environ.get('ANTHROPIC_API_KEY') or os.environ.get('OPENAI_API_KEY'))
provider = 'anthropic' if os.environ.get('ANTHROPIC_API_KEY') else 'openai'
if have_key:
    labels = annotate_codes(adata, 'cell_codebook_idx', provider=provider,
                            tissue='mouse retina', with_literature=True,
                            context_cols=('sample_id',), refine=True)
    attach_labels(adata, 'cell_codebook_idx', labels, key_added='celltype_annot')
    display(labels[['label', 'compartment', 'confidence', 'key_markers']].head(12))
else:
    print('No provider key set; skipping the live LLM call. To run it:')
    print('  export ANTHROPIC_API_KEY=...   # or OPENAI_API_KEY')
    print('  labels = annotate_codes(adata, "cell_codebook_idx", provider="anthropic",')
    print('                          tissue="mouse retina", with_literature=True, refine=True)')
    print('  attach_labels(adata, "cell_codebook_idx", labels, key_added="celltype_annot")')

No provider key set; skipping the live LLM call. To run it:
  export ANTHROPIC_API_KEY=...   # or OPENAI_API_KEY
  labels = annotate_codes(adata, "cell_codebook_idx", provider="anthropic",
                          tissue="mouse retina", with_literature=True, refine=True)
  attach_labels(adata, "cell_codebook_idx", labels, key_added="celltype_annot")

Step 3: name the niches by their cell-type community#

Niches are named from the community of cell types they contain, so annotate cells first, then pass those labels into annotate_niches. This also needs a provider key, so it is guarded the same way.

if have_key and 'celltype_annot' in adata.obs:
    niches = annotate_niches(adata, 'neighborhood_codebook_idx', 'celltype_annot',
                             provider=provider, tissue='mouse retina')
    display(niches.head(12))
else:
    print('Skipping niche labeling (needs a provider key and the cell labels from step 2).')
    print('  niches = annotate_niches(adata, "neighborhood_codebook_idx", "celltype_annot",')
    print('                           provider="anthropic", tissue="mouse retina")')

Skipping niche labeling (needs a provider key and the cell labels from step 2).
  niches = annotate_niches(adata, "neighborhood_codebook_idx", "celltype_annot",
                           provider="anthropic", tissue="mouse retina")

Takeaways#

  • The evidence layer (code_evidence, cluster_codes) is deterministic and runs with no network; it is what makes every label auditable.

  • The LLM layer is optional and grounded: each label must be defensible from the code’s own markers and DEGs, with a literature citation, and low-confidence codes are flagged for expert review.

  • The same workflow labels niches from their cell-type composition. Claude Code and Codex can drive the whole pipeline through the bundled nicheverse-mcp server and the nicheverse-annotate skill.

On this page