API reference #

Returns:

loss – Scalar VQ loss (commitment + codebook MSE if not using EMA + diversity).
quantized – Tensor of shape (B, D, T); carries straight-through gradients.
perplexity – Scalar perplexity exp(H(p)) where p is the batch-average hard-assignment distribution.
encoding_indices – Tensor of shape (B*T, 1) with the assigned code index per row.
Shape
—–
- Input ((B, D, T) with D = embedding_dim, K = num_embeddings.)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.FSQ(embedding_dim, levels=(8, 5, 5, 5), commitment_cost=0.0, diversity_weight=0.0, var_target=1.0)#

Bases: Module

Finite Scalar Quantization (Mentzer et al. 2023): codebook-free discretization.

Projects the latent to len(levels) scalar dims, bounds each with a shifted tanh (using the canonical even-level offset so every level yields that many distinct integer codes), rounds with a straight-through estimator, and encodes a single integer index via mixed-radix. Effective code count is prod(levels); the index is a bijection onto [0, prod(levels)).

Parameters:

embedding_dim (int) – Latent dimensionality; projected to and from len(levels).
levels (tuple[int, ...], default: (8, 5, 5, 5)) – Per-dimension quantization levels (each >= 3; level 2 is degenerate).
commitment_cost (float, default: 0.0) – Accepted for a uniform interface but unused: FSQ has no codebook to commit to; the straight-through round ties the latent to the reconstruction.
diversity_weight (float, default: 0.0) – Optional differentiable range-usage penalty diversity_weight * relu(var_target - var(bounded_z)) (0 disables it; FSQ is collapse-resistant by construction).
var_target (float, default: 1.0) – Optional differentiable range-usage penalty diversity_weight * relu(var_target - var(bounded_z)) (0 disables it; FSQ is collapse-resistant by construction).

Notes

Forward takes (B, D, T) and returns (loss, quantized (B, D, T), perplexity, index (B*T, 1)).

property embedding: Embedding#: Materialize the implicit FSQ codebook (grid points mapped through project_out).

forward(inputs)#

Project to scalar dims, bound and round (straight-through), and encode a mixed-radix index.

The latent is projected to len(levels) scalar channels, each bounded by the shifted tanh and rounded to its integer level with a straight-through estimator, then decoded back to the latent space. The per-dimension integers are packed into one index in [0, prod(levels)) by the mixed-radix _basis. There is no learned codebook, so the loss is zero unless the optional range-usage penalty is enabled.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar loss; 0 unless diversity_weight > 0 adds the range-usage penalty.
quantized (torch.Tensor) – Straight-through quantized latent, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of the batch code-usage histogram.
encoding_indices (torch.Tensor) – Mixed-radix code index per row, shape (B*T, 1), in [0, prod(levels)).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.SoftVQ(num_embeddings, embedding_dim, commitment_cost=0.25, temperature=0.07, diversity_weight=0.01)#

Bases: Module

Soft vector quantizer: differentiable softmax assignment over a learnable codebook.

Parameters:

num_embeddings (int) – Codebook size K and entry dimensionality D.
embedding_dim (int) – Codebook size K and entry dimensionality D.
commitment_cost (float, default: 0.25) – Weight of the soft commitment MSE (quantized toward the input).
temperature (float, default: 0.07) – Softmax temperature on the cosine-similarity logits (lower is sharper).
diversity_weight (float, default: 0.01) – Weight of the marginal-entropy load-balancing term (0 disables it).

Notes

quantized = softmax(cos_sim / temperature) @ codebook is differentiable everywhere (no straight-through estimator); the hard argmax index is returned for annotation. Forward takes (B, D, T) and returns (loss, quantized (B, D, T), perplexity, index (B*T, 1)).

forward(inputs)#

Softly assign the input to the codebook by temperature-scaled cosine similarity.

Returns (loss, quantized, perplexity, index) where quantized is the softmax-weighted codebook average (differentiable, no straight-through) and index is the hard argmax code for annotation.

Return type:: tuple[Tensor, Tensor, Tensor, Tensor]
Parameters:: inputs (Tensor)

Shape#

Input: (B, D, T).
Output: loss (), quantized (B, D, T), perplexity (), index (B*T, 1).

param inputs:
type inputs:: Tensor

class nicheverse.models.RotVQ(num_embeddings, embedding_dim, commitment_cost=0.25, num_householders=8, distance_metric='l2')#

Bases: _VQDelegate, Module

VQ in a learned rotated basis (product of Householder reflections).

Applies an exactly-orthogonal learned rotation to the latent, quantizes with a standard EMA VectorQuantizer in that basis, then applies the exact inverse rotation to the quantized vector. Decorrelates axes to raise codebook utilization; the per-cell code is a single integer. embedding / num_embeddings / embedding_dim / distance_metric delegate to the wrapped quantizer so it satisfies the same contract as VectorQuantizer.

Parameters:

num_embeddings (int) – Codebook size and entry dimensionality of the wrapped EMA quantizer.
embedding_dim (int) – Codebook size and entry dimensionality of the wrapped EMA quantizer.
commitment_cost (float, default: 0.25) – Passed through to the wrapped VectorQuantizer.
distance_metric (str, default: 'l2') – Passed through to the wrapped VectorQuantizer.
num_householders (int, default: 8) – Number of Householder reflections composing the rotation.

forward(inputs)#

Rotate into the learned basis, quantize there, and rotate the result back.

The input is flattened and passed through the composed Householder reflections, the rotated latent is quantized by the wrapped EMA VectorQuantizer (which owns the loss, perplexity, index, and straight-through gradient), and the quantized vector is mapped back with the exact inverse rotation. The code index and loss are those of the wrapped quantizer, so the contract matches VectorQuantizer.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar VQ loss from the wrapped quantizer.
quantized (torch.Tensor) – Inverse-rotated quantized latent, shape (B, D, T); carries straight-through gradients through the wrapped quantizer.
perplexity (torch.Tensor) – Scalar codebook perplexity of the wrapped quantizer.
encoding_indices (torch.Tensor) – Assigned code index per row, shape (B*T, 1).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.QINCoVQ(num_embeddings, embedding_dim, commitment_cost=0.25, num_levels=4, distance_metric='l2')#

Bases: _VQDelegate, Module

QINCo-style conditional residual VQ (Huijben et al. 2024, arXiv:2401.14732).

A stack of num_levels codebooks applied coarse-to-fine on the residual. Level 0 is a plain EMA VectorQuantizer; each deeper level quantizes a query produced by a small MLP conditioned on the residual and the running sum of quantized vectors, so later codes refine earlier ones. The returned single-integer code is the level-0 (coarse) index, which keeps cell_codebook_idx a stable single column; the quantized output is the sum over all levels. embedding / num_embeddings / embedding_dim / distance_metric delegate to level 0.

Parameters:

num_embeddings (int) – Per-level codebook size and entry dimensionality.
embedding_dim (int) – Per-level codebook size and entry dimensionality.
commitment_cost (float, default: 0.25) – Passed to each level’s VectorQuantizer.
distance_metric (str, default: 'l2') – Passed to each level’s VectorQuantizer.
num_levels (int, default: 4) – Number of residual levels (>= 1).

forward(inputs)#

Quantize coarse-to-fine over the residual and sum the per-level quantized vectors.

Level 0 quantizes the input with a plain EMA VectorQuantizer. Each deeper level quantizes a query formed by an MLP conditioned on the running residual and the accumulated quantized sum, so later codes refine earlier ones (the residual is detached before subtracting each level’s output). The output is the sum over all levels; the returned single-integer code and perplexity are level 0’s (coarse) values, keeping cell_codebook_idx a stable single column.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar sum of the per-level VQ losses.
quantized (torch.Tensor) – Sum of the per-level quantized vectors, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of the level-0 codebook.
encoding_indices (torch.Tensor) – Level-0 (coarse) code index per row, shape (B*T, 1).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.ProductVQ(num_embeddings, embedding_dim, commitment_cost=0.25, num_subspaces=4, distance_metric='l2')#

Bases: _VQDelegate, Module

Product quantization: split the latent into num_subspaces chunks, each quantized by its own codebook.

Effective code count is num_embeddings ** num_subspaces at only num_embeddings * embedding_dim parameters, which resists collapse and factorizes a cell into independent programs. The returned single-integer code is subspace 0’s index (kept for a stable cell_codebook_idx); the quantized output is the concatenation over subspaces. embedding / num_embeddings / embedding_dim / distance_metric delegate to subspace 0 (so embedding_dim reports the per-subspace width).

Parameters:

num_embeddings (int) – Codebook size per subspace and total latent dimensionality (must be divisible by num_subspaces).
embedding_dim (int) – Codebook size per subspace and total latent dimensionality (must be divisible by num_subspaces).
commitment_cost (float, default: 0.25) – Passed to each subspace’s VectorQuantizer.
distance_metric (str, default: 'l2') – Passed to each subspace’s VectorQuantizer.
num_subspaces (int, default: 4) – Number of subspaces the latent is split into.

forward(inputs)#

Split the latent into subspaces, quantize each independently, and concatenate.

The channel dimension is split into num_subspaces equal chunks, each quantized by its own EMA VectorQuantizer, and the quantized chunks are concatenated back to the full width. The loss is the sum over subspaces; the returned single-integer code and perplexity are subspace 0’s, keeping cell_codebook_idx a stable single column while the effective code count is num_embeddings ** num_subspaces.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D the full embedding_dim divisible by num_subspaces.

Return type:

Returns:

loss (torch.Tensor) – Scalar sum of the per-subspace VQ losses.
quantized (torch.Tensor) – Concatenation of the per-subspace quantized latents, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of subspace 0’s codebook.
encoding_indices (torch.Tensor) – Subspace-0 code index per row, shape (B*T, 1).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.ResidualVQ(num_embeddings, embedding_dim, commitment_cost=0.25, num_stages=4, distance_metric='l2')#

Bases: _VQDelegate, Module

Residual vector quantization: a stack of EMA VectorQuantizer stages.

Stage 0 quantizes the input; each later stage quantizes the residual left by the running sum of quantized vectors (the residual is detached before subtracting each stage’s output, exactly as in QINCoVQ), so deeper stages refine coarser ones. The output is the sum over stages; the returned single-integer code and perplexity are stage 0’s (coarse) values, keeping cell_codebook_idx a stable single column. embedding / num_embeddings / embedding_dim / distance_metric delegate to stage 0. Unlike QINCoVQ there is no conditioning MLP: each stage is a plain residual codebook.

Parameters:

num_embeddings (int) – Per-stage codebook size and entry dimensionality.
embedding_dim (int) – Per-stage codebook size and entry dimensionality.
commitment_cost (float, default: 0.25) – Passed to each stage’s VectorQuantizer.
distance_metric (str, default: 'l2') – Passed to each stage’s VectorQuantizer.
num_stages (int, default: 4) – Number of residual stages (>= 1).

forward(inputs)#

Quantize the residual stage by stage and sum the per-stage quantized vectors.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar sum of the per-stage VQ losses.
quantized (torch.Tensor) – Sum of the per-stage quantized vectors, shape (B, D, T); carries straight-through gradients.
perplexity (torch.Tensor) – Scalar perplexity of stage 0’s codebook.
encoding_indices (torch.Tensor) – Stage-0 (coarse) code index per row, shape (B*T, 1).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.LFQ(embedding_dim, commitment_cost=0.25, entropy_weight=0.1, commitment_weight=None)#

Bases: Module

Lookup-Free Quantization (Yu et al., MAGVIT-v2, ICLR 2024, arXiv:2310.05737).

There is no learned codebook: each latent dimension becomes one bit via its sign, so the implicit codebook is {-1, +1}**embedding_dim and the effective code count is 2**embedding_dim. The integer code is the bit-packed sign pattern. A commitment term ties the latent to its binarization and an entropy term (marginal minus conditional bit entropy) prevents dead bits.

Parameters:

embedding_dim (int) – Latent dimensionality; also the number of bits (code count 2**embedding_dim).
commitment_cost (float, default: 0.25) – Default weight of the commitment term (used when commitment_weight is None).
entropy_weight (float, default: 0.1) – Weight of the bit-entropy regularizer.
commitment_weight (Optional[float], default: None) – Explicit commitment weight; falls back to commitment_cost when None.

Notes

embedding_dim is the working (bit) dimensionality and num_embeddings is the nominal code count 2**embedding_dim. To stay within int64, the packed index uses at most 62 bits (so the index is exact for embedding_dim <= 62 and always in [0, num_embeddings)). Forward takes (B, D, T) and returns (loss, quantized (B, D, T), perplexity, index (B*T, 1)).

forward(inputs)#

Binarize each latent dimension by its sign and pack the bits into one integer code.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar commitment + entropy loss.
quantized (torch.Tensor) – Straight-through binarized latent (+/-1 per dim), shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of the (capped) batch code histogram.
encoding_indices (torch.Tensor) – Bit-packed code index per row, shape (B*T, 1), in [0, 2**D).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.BSQ(embedding_dim, commitment_cost=0.25, gamma=1.0)#

Bases: Module

Binary Spherical Quantization (Zhao et al., ICML 2024, arXiv:2406.07548).

The latent is L2-normalized onto the unit sphere and then binarized per dimension to +/- 1/sqrt(D), so codes lie on hypercube vertices on the sphere and the effective code count is 2**embedding_dim. The integer code is the bit-packed sign pattern. A marginal bit-entropy term encourages balanced bit usage.

Parameters:

embedding_dim (int) – Latent dimensionality; also the number of bits (code count 2**embedding_dim).
commitment_cost (float, default: 0.25) – Accepted for a uniform interface; BSQ’s loss is the entropy regularizer only.
gamma (float, default: 1.0) – Temperature of the soft-sign used in the bit-entropy regularizer.

Notes

embedding_dim is the working (bit) dimensionality and num_embeddings is the nominal code count 2**embedding_dim; the packed index uses at most 62 bits so it stays in [0, num_embeddings) within int64. Forward takes (B, D, T) and returns (loss, quantized (B, D, T), perplexity, index (B*T, 1)).

forward(inputs)#

Project onto the unit sphere, binarize each dimension, and pack the bits into a code.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar marginal bit-entropy regularizer.
quantized (torch.Tensor) – Straight-through spherical binarized latent, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of the (capped) batch code histogram.
encoding_indices (torch.Tensor) – Bit-packed code index per row, shape (B*T, 1), in [0, 2**D).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.ResidualFSQ(embedding_dim, levels=(8, 5, 5, 5), num_stages=3, commitment_cost=0.0, diversity_weight=0.0)#

Bases: Module

Residual finite scalar quantization (arXiv:2508.15860): multi-stage FSQ.

Stacks num_stages codebook-free FSQ layers, each quantizing the residual left by the previous stages (residual detached before subtracting), so the stages form a hierarchical, collapse-resistant code. The output is the sum over stages; the returned single-integer code and perplexity are stage 0’s, keeping cell_codebook_idx a stable single column. num_embeddings reports the per-stage code count prod(levels) and embedding_dim the working dimension.

Parameters:

embedding_dim (int) – Latent dimensionality (each stage projects to and from len(levels)).
levels (tuple[int, ...], default: (8, 5, 5, 5)) – Per-dimension quantization levels shared by every stage (each >= 3).
num_stages (int, default: 3) – Number of residual FSQ stages (>= 1).
commitment_cost (float, default: 0.0) – Passed to each stage’s FSQ (FSQ is codebook-free; both default off).
diversity_weight (float, default: 0.0) – Passed to each stage’s FSQ (FSQ is codebook-free; both default off).

forward(inputs)#

Quantize the residual with a fresh FSQ at each stage and sum the quantized outputs.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D = embedding_dim.

Return type:

Returns:

loss (torch.Tensor) – Scalar sum of the per-stage FSQ losses (0 unless a stage penalty is on).
quantized (torch.Tensor) – Sum of the per-stage quantized vectors, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of stage 0.
encoding_indices (torch.Tensor) – Stage-0 (coarse) mixed-radix code index per row, shape (B*T, 1), in [0, prod(levels)).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

class nicheverse.models.GroupedResidualVQ(num_embeddings, embedding_dim, commitment_cost=0.25, num_groups=2, num_stages=2, distance_metric='l2')#

Bases: _VQDelegate, Module

Grouped residual VQ: split the channels into groups, an ResidualVQ per group.

The channel dimension is split into num_groups equal chunks; each chunk is quantized by its own residual VQ, and the quantized chunks are concatenated back to the full width (Yang et al., HiFi-Codec, arXiv:2305.02765). The loss is the sum over groups; the returned single-integer code and perplexity are group 0’s stage-0 values, keeping cell_codebook_idx a stable single column. embedding / num_embeddings / embedding_dim / distance_metric delegate to group 0 (so embedding_dim reports the per-group width).

Parameters:

num_embeddings (int) – Per-stage codebook size and total latent dimensionality (must be divisible by num_groups).
embedding_dim (int) – Per-stage codebook size and total latent dimensionality (must be divisible by num_groups).
commitment_cost (float, default: 0.25) – Passed to each group’s ResidualVQ.
distance_metric (str, default: 'l2') – Passed to each group’s ResidualVQ.
num_groups (int, default: 2) – Number of channel groups the latent is split into.
num_stages (int, default: 2) – Number of residual stages within each group’s ResidualVQ.

forward(inputs)#

Split into channel groups, run an independent residual VQ per group, and concatenate.

Parameters:

inputs (torch.Tensor) – Encoded latent, shape (B, D, T) (channels-first; the model uses T=1), with D divisible by num_groups.

Return type:

Returns:

loss (torch.Tensor) – Scalar sum of the per-group residual VQ losses.
quantized (torch.Tensor) – Concatenation of the per-group quantized latents, shape (B, D, T).
perplexity (torch.Tensor) – Scalar perplexity of group 0’s stage-0 codebook.
encoding_indices (torch.Tensor) – Group-0 stage-0 code index per row, shape (B*T, 1).
Shape
—–
- Input ((B, D, T).)
- Output (loss (), quantized (B, D, T), perplexity) – (), encoding_indices (B*T, 1).

nicheverse.models.build_quantizer(name, *, num_embeddings, embedding_dim, commitment_cost, distance_metric='l2', **kwargs)#

Construct a quantizer by registry name (see _QUANTIZERS for the full set).

build_quantizer("vq", ...) returns exactly the VectorQuantizer the model builds by default, so the default path is bit-identical. The codebook-free variants (fsq, lfq, bsq, residual_fsq) derive their code count implicitly and ignore num_embeddings (a warning is logged if the two disagree).

Parameters:

name (str)
num_embeddings (int)
embedding_dim (int)
commitment_cost (float)
distance_metric (str, default: 'l2')

nicheverse.models.register_quantizer(name)#

Class decorator registering a quantizer under name for build_quantizer().

Parameters:: name (str)

Encoders#

Encoder backbones map each cell’s gene vector to a latent, chosen with ModelConfig.encoder_type and configured with ModelConfig.encoder_kwargs. All share the builder signature (in_dim, out_dim, hidden, dropout, **kwargs) and map (batch, in_dim) to (batch, out_dim).

`encoder_type`	Backbone	`encoder_kwargs`
`mlp_deep` (default)	Deep, wide SwiGLU pre-norm residual MLP (no per-gene numerical embedding)	`ffn_mult`, `n_blocks`
`mlp`	Plain MLP: Linear, BatchNorm, ReLU, Dropout	(none)
`mlp_plr`	MLP with per-gene periodic numerical embeddings and a SwiGLU trunk	`n_freq`, `num_emb_dim`, `n_blocks`, `freq_init_sigma`, `ffn_mult`
`residual_mlp`	Pre-activation residual MLP blocks	(none)
`transformer`	Gene-patch Transformer	`patch_size`, `num_heads`, `num_layers`
`cnn`	1D convolutional network over the gene vector	`channels`, `kernel_sizes`, `use_batch_norm`
`fast_cnn`	Lightweight 1D CNN	`channels`, `kernel_sizes`, `use_batch_norm`
`deep_cnn`	Deep residual 1D CNN with squeeze-and-excitation and multiscale fusion	`base_channels`, `depth_multiplier`, `n_stages`, `res_blocks_per_stage`, `kernel_sizes`, `use_se`, `use_multiscale`
`gnn`	Graph-attention message passing over gene-patch tokens	`patch_size`, `num_layers`
`diffusion`	U-Net denoiser encoder	`time_emb_dim`, `n_timesteps`
`dit`	Diffusion Transformer with AdaLN conditioning	`patch_size`, `num_heads`, `num_layers`, `time_emb_dim`, `n_timesteps`
`set_transformer`	Set Transformer (ISAB, PMA pooling) over gene-patch tokens	`patch_size`, `num_heads`, `num_inds`, `num_isab`
`perceiver_io`	Perceiver IO latent cross-attention	`patch_size`, `num_heads`, `num_latents`, `num_self_layers`
`soft_moe`	Soft Mixture-of-Experts routing	`patch_size`, `num_heads`, `num_experts`, `slots_per_expert`
`ft_transformer`	FT-Transformer with per-gene feature tokens (compute-prohibitive at cohort scale)	`num_heads`, `num_layers`

nicheverse.models.build_encoder(name, *, in_dim, out_dim, hidden, dropout=0.2, **kwargs)#

Build an encoder by registry name. "mlp" (default) is the released encoder.

Extra keyword arguments are forwarded to the encoder builder (e.g. patch_size / num_heads / num_layers for "transformer").

Parameters:

name (str)
in_dim (int)
out_dim (int)
hidden (Sequence[int])
dropout (float, default: 0.2)
kwargs (Any)

Return type:

Module

nicheverse.models.register_encoder(name)#

Register an encoder builder (callable (in_dim, out_dim, hidden, dropout) -> nn.Module).

Parameters:: name (str)

class nicheverse.models.ResidualMLP(in_dim, out_dim, hidden, dropout=0.2)#

Bases: Module

Pre-activation residual MLP encoder (LayerNorm, Linear, GELU, Dropout, Linear + skip).

Parameters:

in_dim (int) – Input and output dimensionality.
out_dim (int) – Input and output dimensionality.
hidden – Per-block intermediate widths; the residual stream width is hidden[0].
dropout (float, default: 0.2) – Dropout probability inside each residual block.

forward(x)#

Encode a gene vector to a latent through the residual stack.

Return type:: Tensor
Parameters:: x (Tensor)

Shape#

Input: (B, in_dim) per-cell (normalized) gene vector.
Output: (B, out_dim) latent embedding.

param x:
type x:: Tensor

class nicheverse.models.TransformerEncoder(in_dim, out_dim, hidden, dropout=0.2, patch_size=16, num_heads=4, num_layers=2)#

Bases: Module

Gene-patch Transformer encoder: patchify the gene vector, self-attend, mean-pool to a latent.

Models combinatorial (nonlinear) gene interactions an MLP only mixes additively. Operates on the existing per-cell gene vector, no extra inputs.

Parameters:

in_dim (int) – Input gene count and output latent dimensionality.
out_dim (int) – Input gene count and output latent dimensionality.
hidden (Sequence[int]) – The model width is hidden[0].
dropout (float, default: 0.2) – Attention and feed-forward dropout.
patch_size (int, default: 16) – Genes per patch token.
num_heads (int, default: 4) – Requested attention heads (reduced to the largest divisor of the width).
num_layers (int, default: 2) – Number of Transformer encoder layers.

forward(x)#

Patchify, self-attend over gene patches, mean-pool, and project to a latent.

Return type:: Tensor
Parameters:: x (Tensor)

Shape#

Input: (B, in_dim) gene vector (right-padded to a multiple of patch_size).
Intermediate: (B, n_patches, hidden[0]) patch tokens.
Output: (B, out_dim) latent embedding (mean over patch tokens).

param x:
type x:: Tensor

Molecule-set model#

An alternative input modality: instead of a per-cell aggregated gene vector, each cell is encoded as the set of its transcript molecules (one token per molecule, carrying gene identity and subcellular offset) through a permutation-invariant, mask-aware Set Transformer. :class:~nicheverse.models.MoleculeSetVQVAE mirrors :class:~nicheverse.HierarchicalVQVAE (cell codebook, optional neighborhood codebook, cross-attention fusion) and reuses the released :class:~nicheverse.VectorQuantizer, so the codebook behavior and emitted code indices are unchanged.

class nicheverse.models.MoleculeSetEncoder(out_dim=64, width=128, n_genes=366, num_freqs=12, num_heads=4, num_inds=16, num_isab=2, dropout=0.1, max_micron=7.0)#

Bases: Module

Set Transformer++ over a cell’s molecule set -> per-cell latent.

Each molecule token is gene_embedding(gene_id) + posMLP(Fourier(dx, dy)); masked ISAB blocks mix the set; the set is pooled by concatenating masked max, masked mean, and a masked PMA seed readout, then projected to out_dim. PMA alone collapses (its single seed query converges to a near-constant readout, so the codebook EMA chases it into one cluster); the max/mean statistics preserve each cell’s molecular-composition magnitude and spread, giving the quantizer real structure. For an all-padded (empty) cell the masked max and masked mean are already 0, but the PMA seed reads a garbage padded token, so it is zeroed for empty rows, giving a fully zero pooled embedding. Permutation invariant and mask aware (True = real molecule; gene_id == n_genes is the padding token).

Parameters:

out_dim (int, default: 64)
width (int, default: 128)
n_genes (int, default: 366)
num_freqs (int, default: 12)
num_heads (int, default: 4)
num_inds (int, default: 16)
num_isab (int, default: 2)
dropout (float, default: 0.1)
max_micron (float, default: 7.0)

forward(gene, coords, mask)#

Encode molecules gene:(B,M) long, coords:(B,M,2), mask:(B,M) bool to a latent (B, out_dim).

Parameters:

gene (Tensor)
coords (Tensor)
mask (Tensor)

Return type:

Tensor

class nicheverse.models.MoleculeSetVQVAE(n_genes=366, cell_embedding_dim=64, cell_num_embeddings=256, neighborhood_embedding_dim=256, neighborhood_num_embeddings=32, commitment_cost=0.25, hidden=(256, 128), use_neighborhood=True, use_cross_attention=True, cross_attention_weight=0.5, enc_width=128, enc_isab=2, enc_heads=4, enc_inds=16, enc_freqs=12)#

Bases: Module

Hierarchical VQ-VAE with a molecule-set cell branch and an optional aggregated kNN neighborhood branch, fused by cross-attention (mirrors HierarchicalVQVAE). Emits the same cell/neighborhood code indices so downstream annotation is unchanged.

Parameters:

n_genes (int, default: 366)
cell_embedding_dim (int, default: 64)
cell_num_embeddings (int, default: 256)
neighborhood_embedding_dim (int, default: 256)
neighborhood_num_embeddings (int, default: 32)
commitment_cost (float, default: 0.25)
use_neighborhood (bool, default: True)
use_cross_attention (bool, default: True)
cross_attention_weight (float, default: 0.5)
enc_width (int, default: 128)
enc_isab (int, default: 2)
enc_heads (int, default: 4)
enc_inds (int, default: 16)
enc_freqs (int, default: 12)

forward(gene, coords, mask, neigh_features=None)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Data#

Reading imaging spatial transcriptomics into an AnnData and building the spatial neighbor graph the model consumes.

class nicheverse.SpatialDataset(cell_features, spatial_coords, sample_ids, k_neighbors=20, neighborhood_aggregation='weighted_mean', spatial_graph='knn_radius', radius=50.0, bandwidth=None, agg_chunk=100000, recon_target=None)#

Bases: Dataset

Cell + neighborhood-aggregated feature pairs, computed per sample.

Parameters:

cell_features (Any) – (n_cells, n_features) dense ndarray or scipy sparse matrix.
spatial_coords (Any) – (n_cells, 2) physical coordinates in microns.
sample_ids (Any) – (n_cells,) string label per cell. The spatial graph is computed within sample only; cells from different samples never become neighbors.
k_neighbors (int, default: 20) – Neighbors per cell (excluding self) for the "knn" graph, and the cap on neighbors kept for "radius" / "delaunay". If a sample has fewer than k_neighbors + 1 cells, all available cells are used. A single-cell sample falls back to duplicating the cell’s own vector.
neighborhood_aggregation (str, default: 'weighted_mean') – One of {"mean", "weighted_mean", "max", "gaussian", "inverse_square"}. "weighted_mean" uses 1 / (distance + epsilon) weights, "inverse_square" uses 1 / (distance ** 2 + epsilon), and "gaussian" uses exp(-distance ** 2 / (2 * bandwidth ** 2)).
spatial_graph (str, default: 'knn_radius') – One of {"knn", "radius", "delaunay", "alpha_complex"}. "knn" (default) keeps the k_neighbors nearest cells. "radius" keeps every cell within radius microns (capped at k_neighbors). "delaunay" uses the Delaunay triangulation adjacency; "alpha_complex" additionally prunes long triangulation edges. Both fall back to a "knn" graph for samples with fewer than three cells or degenerate (collinear) coordinates.
radius (float | None, default: 50.0) – Neighborhood radius in microns; required when spatial_graph="radius".
bandwidth (Optional[float], default: None) – Gaussian kernel bandwidth in microns; required when neighborhood_aggregation="gaussian".
agg_chunk (int, default: 100000) – Cells per vectorized aggregation chunk. Bounds peak memory. Defaults to 100,000.
recon_target (Any | None)

Raises:

ValueError – If neighborhood_aggregation / spatial_graph is unsupported, if input shapes are inconsistent, if k_neighbors is non-positive, if spatial_graph="radius" without a positive radius, or if neighborhood_aggregation="gaussian" without a positive bandwidth.

Notes

Determinism. The neighbor search produces a deterministic ordering given deterministic input: the index into sample_ids defines the row ordering inside each per-sample subset, np.unique returns samples in sorted order, and ties at the boundary distance are broken in input-row order.

Parameters:

recon_target (Optional[Any], default: None)
cell_features (Any)
spatial_coords (Any)
sample_ids (Any)
k_neighbors (int)
neighborhood_aggregation (str)
spatial_graph (str)
radius (float | None)
bandwidth (float | None)
agg_chunk (int)

classmethod from_anndata(adata, sample_col='sample_id', spatial_key='spatial', x_col=None, y_col=None, coord_scale=1.0, transcript_context_key=None, recon_target_layer=None, **kwargs)#

Build a dataset from an AnnData of any imaging spatial-transcriptomics platform.

Coordinates are taken from obsm[spatial_key] when present, otherwise from the obs columns x_col / y_col, and multiplied by coord_scale to convert to microns (e.g. CosMx pixels use coord_scale=0.12028). If sample_col is absent, all cells are treated as a single sample.

Parameters:

adata (Any) – AnnData with a cell-by-gene X and spatial coordinates.
sample_col (str, default: 'sample_id') – obs column naming the per-sample (or per-FOV) unit.
spatial_key (str, default: 'spatial') – obsm key holding an (n_cells, 2) coordinate array.
x_col (Optional[str], default: None) – Fallback obs coordinate columns used when obsm[spatial_key] is absent.
y_col (Optional[str], default: None) – Fallback obs coordinate columns used when obsm[spatial_key] is absent.
coord_scale (float, default: 1.0) – Multiplier converting the stored coordinates to microns.
transcript_context_key (Optional[str], default: None) – If given, concatenate obsm[transcript_context_key] (segmentation-free transcript-context features from transcript_context()) onto X to form the cell features. The model input_dim must then be X.shape[1] + obsm[key].shape[1].
recon_target_layer (Optional[str], default: None) – If given, take the per-cell raw-count reconstruction target from adata.layers[recon_target_layer] (shape must match X). Used by the count-recon cell modes (cell_recon in {"nb","poisson","both"}) so the encoder can see log1p X while the likelihood models the raw counts. None (default) leaves the dataset with no separate recon target (MSE reconstructs the input).
**kwargs (Any) – Forwarded to SpatialDataset (k_neighbors, spatial_graph, …).

Return type:

SpatialDataset

The spatial graph mode is chosen with SpatialDataset.from_anndata(..., spatial_graph=...) (mirrored by TrainConfig.spatial_graph): knn_radius (default, a k-NN graph whose edges are capped at radius microns, default 50), knn (pure k-NN, no cap), radius, delaunay, alpha_complex, gabriel, and rng. The graph is always built per sample, so cells from different samples never link. Passing transcript_context_key concatenates a segmentation-free molecular field (see :func:~nicheverse.data.transcript_context) onto the counts, doubling the model input dimension.

nicheverse.read_spatial(source, sample_col='sample_id', spatial_key='spatial', x_col=None, y_col=None, coord_scale=1.0, copy=False)#

Standardize an AnnData from any imaging spatial-transcriptomics platform.

Accepts an AnnData or an .h5ad path and returns an AnnData ready for train_model(), Trainer, or predict_codes(). It ensures obsm[spatial_key] holds micron coordinates (building it from obs[x_col] / obs[y_col] and scaling by coord_scale when needed) and that obs[sample_col] exists.

Parameters:

source (Any) – An AnnData or a path to an .h5ad file.
sample_col (str, default: 'sample_id') – obs column naming the per-sample (or per-FOV) unit; created as a single sample if absent.
spatial_key (str, default: 'spatial') – obsm key for the (n_cells, 2) micron coordinates.
x_col (Optional[str], default: None) – obs coordinate columns used to build obsm[spatial_key] when absent.
y_col (Optional[str], default: None) – obs coordinate columns used to build obsm[spatial_key] when absent.
coord_scale (float, default: 1.0) – Multiplier converting coordinates to microns (e.g. CosMx pixels use 0.12028).
copy (bool, default: False) – If True and source is already an AnnData, operate on a copy.

Returns:

The standardized AnnData.

Return type:

anndata.AnnData

nicheverse.read_xenium(run_dir, sample_id=None, drop_controls=True)#

Load one Xenium output directory into an anndata.AnnData.

Parameters:

run_dir (str | Path) – Path to a Xenium output directory containing cell_feature_matrix.h5 and cells.parquet (or their fallbacks).
sample_id (Optional[str], default: None) – Label written to obs['sample_id'] and appended to every barcode as cell_id__sample_id. Defaults to Path(run_dir).name.
drop_controls (bool, default: True) – Drop control / blank / unassigned / negative-control / codeword probes from var (case-insensitive match).

Returns:

Cells x genes matrix with obs['sample_id'], obsm['spatial'] (microns), and obs joined with the contents of cells.parquet.

Return type:

AnnData

Raises:

FileNotFoundError – If run_dir does not exist or is missing required files.
ValueError – If cells.parquet is missing the required x_centroid / y_centroid columns.

nicheverse.read_xenium_cohort(run_dirs, sample_ids=None, drop_controls=True, manifest=None)#

Merge multiple Xenium runs into a single anndata.AnnData.

Parameters:

run_dirs (Iterable[str | Path] | str | Path) – One Xenium output directory, an iterable of such directories, or unused (pass []) when manifest is provided.
sample_ids (Optional[Iterable[str]], default: None) – Optional iterable of sample IDs matching run_dirs length. Defaults to Path(p).name for each run.
drop_controls (bool, default: True) – Drop control / blank / unassigned / codeword probes (see load_xenium_run()).
manifest (Union[str, Path, None], default: None) – CSV with columns run_dir and sample_id. If provided, takes precedence over run_dirs / sample_ids.

Returns:

Cohort AnnData restricted to genes shared across all runs, with a categorical obs['sample_id'].

Return type:

AnnData

Raises:

ValueError – If run_dirs is empty and no manifest is given, lengths mismatch, or no genes are shared across the runs.

nicheverse.load_xenium_run(run_dir, sample_id=None, drop_controls=True)#

Load one Xenium output directory into an anndata.AnnData.

Parameters:

run_dir (str | Path) – Path to a Xenium output directory containing cell_feature_matrix.h5 and cells.parquet (or their fallbacks).
sample_id (Optional[str], default: None) – Label written to obs['sample_id'] and appended to every barcode as cell_id__sample_id. Defaults to Path(run_dir).name.
drop_controls (bool, default: True) – Drop control / blank / unassigned / negative-control / codeword probes from var (case-insensitive match).

Returns:

Cells x genes matrix with obs['sample_id'], obsm['spatial'] (microns), and obs joined with the contents of cells.parquet.

Return type:

AnnData

Raises:

FileNotFoundError – If run_dir does not exist or is missing required files.
ValueError – If cells.parquet is missing the required x_centroid / y_centroid columns.

nicheverse.load_xenium_cohort(run_dirs, sample_ids=None, drop_controls=True, manifest=None)#

Merge multiple Xenium runs into a single anndata.AnnData.

Parameters:

run_dirs (Iterable[str | Path] | str | Path) – One Xenium output directory, an iterable of such directories, or unused (pass []) when manifest is provided.
sample_ids (Optional[Iterable[str]], default: None) – Optional iterable of sample IDs matching run_dirs length. Defaults to Path(p).name for each run.
drop_controls (bool, default: True) – Drop control / blank / unassigned / codeword probes (see load_xenium_run()).
manifest (Union[str, Path, None], default: None) – CSV with columns run_dir and sample_id. If provided, takes precedence over run_dirs / sample_ids.

Returns:

Cohort AnnData restricted to genes shared across all runs, with a categorical obs['sample_id'].

Return type:

AnnData

Raises:

ValueError – If run_dirs is empty and no manifest is given, lengths mismatch, or no genes are shared across the runs.

nicheverse.spatial_neighbors(adata, k_neighbors=20, neighborhood_aggregation='weighted_mean', spatial_graph='knn_radius', radius=50.0, bandwidth=None, sample_col='sample_id', key_added='neighborhood_features', copy=False)#

Compute per-sample aggregated neighborhood features and store them in obsm.

The neighborhood feature for each cell is the concatenation of its own expression and an aggregation over its spatial neighbors, computed within each sample so cross-sample edges never form (see SpatialDataset).

Parameters:

adata (AnnData) – AnnData with obsm['spatial'] (microns) and obs[sample_col].
k_neighbors (int, default: 20) – Forwarded to SpatialDataset.
neighborhood_aggregation (str, default: 'weighted_mean') – Forwarded to SpatialDataset.
spatial_graph (str, default: 'knn_radius') – Forwarded to SpatialDataset.
radius (float | None, default: 50.0) – Forwarded to SpatialDataset.
bandwidth (Optional[float], default: None) – Forwarded to SpatialDataset.
sample_col (str, default: 'sample_id') – obs column partitioning cells for the per-sample graph.
key_added (str, default: 'neighborhood_features') – obsm key to write the (n_cells, 2 * n_genes) feature matrix to.
copy (bool, default: False) – If True, operate on and return a copy; otherwise write in place and return the (n_cells, 2 * n_genes) ndarray.

Returns:

The mutated (or copied) AnnData when copy=True, else the feature matrix ndarray.

Return type:

AnnData or numpy.ndarray

nicheverse.data.transcript_context(adata, transcripts, radius=7.0, sample_col='sample_id', platform='xenium', x_col=None, y_col=None, feature_col=None, control_pattern=None, key_added='transcript_context', copy=False)#

Compute the per-cell local molecular field and store it in obsm.

Parameters:

adata (AnnData) – AnnData with obsm['spatial'] (microns) and obs[sample_col].
transcripts (dict | str | Path) – Either a mapping {sample_id: parquet_path} (one molecule table per sample) or a single path used for every sample (single-sample runs).
radius (float, default: 7.0) – Micron radius of the molecular field around each centroid.
platform (str, default: 'xenium') – Molecule-table convention: "xenium" (x_location/y_location/ feature_name), "cosmx" (x_global_px/y_global_px/target), or "merfish" (global_x/global_y/gene). Sets column and control-probe defaults.
x_col (Optional[str], default: None) – Molecule table column names; override the platform defaults when given.
y_col (Optional[str], default: None) – Molecule table column names; override the platform defaults when given.
feature_col (Optional[str], default: None) – Molecule table column names; override the platform defaults when given.
control_pattern (Optional[Pattern], default: None) – Compiled regex of control/blank probe names to drop; defaults to the platform pattern.
key_added (str, default: 'transcript_context') – obsm key for the (n_cells, n_genes) log1p count matrix.
copy (bool, default: False) – If True, operate on and return a copy; else write in place and return the feature matrix.
sample_col (str, default: 'sample_id')

Raises:

ValueError – If obsm['spatial'] / obs[sample_col] is missing, or a sample has no transcripts path.

Return type:

AnnData | ndarray

nicheverse.attach_codes(adata, cell_idx, neigh_idx, cell_emb=None, neigh_emb=None)#

Write codebook indices and (optionally) embeddings back into AnnData.

Indices are stored as the narrowest integer dtype that fits their range, saving disk space for the typical case where the codebook has at most a few hundred entries.

Parameters:

adata (AnnData) – AnnData with adata.n_obs cells; mutated in place.
cell_idx (ndarray) – (n_obs,) integer arrays of code assignments.
neigh_idx (ndarray) – (n_obs,) integer arrays of code assignments.
cell_emb (Optional[ndarray], default: None) – Optional (n_obs, D) continuous embeddings written to adata.obsm.
neigh_emb (Optional[ndarray], default: None) – Optional (n_obs, D) continuous embeddings written to adata.obsm.

Returns:

The same adata with new columns and (optionally) obsm matrices.

Return type:

AnnData

Raises:

ValueError – If any input array length differs from adata.n_obs.

nicheverse.attach_codes_to_adata(adata, cell_idx, neigh_idx, cell_emb=None, neigh_emb=None)#

Write codebook indices and (optionally) embeddings back into AnnData.

Indices are stored as the narrowest integer dtype that fits their range, saving disk space for the typical case where the codebook has at most a few hundred entries.

Parameters:

adata (AnnData) – AnnData with adata.n_obs cells; mutated in place.
cell_idx (ndarray) – (n_obs,) integer arrays of code assignments.
neigh_idx (ndarray) – (n_obs,) integer arrays of code assignments.
cell_emb (Optional[ndarray], default: None) – Optional (n_obs, D) continuous embeddings written to adata.obsm.
neigh_emb (Optional[ndarray], default: None) – Optional (n_obs, D) continuous embeddings written to adata.obsm.

Returns:

The same adata with new columns and (optionally) obsm matrices.

Return type:

AnnData

Raises:

ValueError – If any input array length differs from adata.n_obs.

The subcellular molecule-set representation (a per-cell transcript point cloud) is loaded by :class:~nicheverse.data.MoleculeSetDataset and consumed by :class:~nicheverse.models.MoleculeSetVQVAE.

class nicheverse.data.MoleculeSetDataset(obs_names, shard_dir, spatial_coords=None, sample_ids=None, expr_lognorm=None, k_neighbors=20, neighborhood_aggregation='weighted_mean', with_neighborhood=True)#

Bases: Dataset

Per-cell molecule sets aligned to obs_names (+ optional kNN neighborhood).

Parameters:

obs_names – Cell identifiers defining the row order (aligns shards to the model adata).
shard_dir – Directory of *.npz molecule-set shards.
spatial_coords (default: None) – Passed to SpatialDataset to build the neighborhood branch when with_neighborhood is True.
sample_ids (default: None) – Passed to SpatialDataset to build the neighborhood branch when with_neighborhood is True.
expr_lognorm (default: None) – Passed to SpatialDataset to build the neighborhood branch when with_neighborhood is True.
k_neighbors (int, default: 20) – Forwarded to SpatialDataset.
neighborhood_aggregation (str, default: 'weighted_mean') – Forwarded to SpatialDataset.
with_neighborhood (bool, default: True) – If True, attach (n, 2*n_genes) aggregated neighborhood features.

Training#

class nicheverse.Trainer(config=None, model_config=None)#

Bases: object

High-level training and inference wrapper (nanoGPT / Lightning style).

Parameters:

config (default: None) – A TrainConfig; defaults to the reproducible production settings.
model_config (default: None) – Optional ModelConfig; if omitted, fit builds one from the AnnData gene panel.

Examples

>>> trainer = Trainer(TrainConfig(num_epochs=10))
>>> model, adata = trainer.fit(adata, "checkpoints/")
>>> annotated = trainer.predict(new_adata, "checkpoints/hierarchical_vqvae_checkpoint.pt")

fit(adata, checkpoint_dir, model_config=None, sample_col='sample_id', device=None)#: Train and persist the model; returns (model, annotated_adata).

predict(adata, checkpoint=None, **kwargs)#: Assign codes to adata using a checkpoint path or the fitted model.

class nicheverse.TrainConfig(num_epochs=300, batch_size=32768, scale_lr_with_batch=True, learning_rate=0.0003, weight_decay=0.01, lr_schedule='plateau', warmup_steps=0, warmup_frac=0.03, min_lr=0.0, decoupled_weight_decay=True, spatial_loss_type='none', spatial_loss_weight=0.0, spatial_loss_k=6, cell_weight=1.0, neighborhood_weight=1.0, k_neighbors=20, neighborhood_aggregation='weighted_mean', spatial_graph='knn_radius', radius=50.0, bandwidth=None, normalize=True, log1p=True, val_fraction=0.0, early_stopping_patience=None, grad_clip=None, amp=False, resume_from=None, save_best=True, seed=9, log_every=10, deterministic=True, num_workers=0, ddp=False, compile_model=False, pin_memory=None, persistent_workers=None, prefetch_factor=None, device_resident=True, device_resident_mem_fraction=0.9)#

Bases: object

Hyperparameters for train_model().

Parameters:

num_epochs (int, default: 300) – Number of full passes over the dataset.
batch_size (int | str, default: 32768) – Cells per mini-batch. An int (default 32768) or the string "auto" to resolve it at train time from the panel size via auto_batch_size(). Only "auto" triggers the adaptive path; any int is used verbatim. The default was raised from 2048 to 32768 following the large-batch literature (McCandlish 2018 critical batch size; Goyal 2017 linear-scaling + warmup; Smith 2018): at cohort scale (millions of cells) a much larger batch improves gradient-estimate quality and hardware utilization with no loss of final quality when paired with LR warmup (see lr_schedule='warmup_cosine' and warmup_frac). A batch of 32768 was measured at ~45 GB resident on an 80 GB A100, so it fits the device_resident path with headroom; the power-of-two floor logic in auto_batch_size() and the device_resident fit-check both hold for 32768. The default 32768 is the released b32k trajectory; pass a smaller int (e.g. 2048) only for small-GPU or quick runs.
scale_lr_with_batch (bool, default: True) – When batch_size="auto" and this is True (default), scale the learning rate by sqrt(effective_batch / 2048) so a larger resolved batch keeps a comparable update magnitude. No effect for an int batch.
learning_rate (float, default: 0.0003) – Adam learning rate. Reduced on plateau by factor=0.5 with patience=5 epochs (monitors the validation loss when val_fraction > 0, else the training loss).
weight_decay (float, default: 0.01) – AdamW decoupled weight decay coefficient (Loshchilov and Hutter 2019). 0.01 (default) is the AdamW reference value and is conservative for a VQ tokenizer trained on millions of cells (large data, low overfitting risk). Applied ONLY when decoupled_weight_decay is True (the default), and then only to the 2-D-or-wider weight matrices of nn.Linear / nn.Conv* and the attention QKV projection; every other parameter (both VQ codebook embeddings, the molecule-set gene embedding, all nn.LayerNorm / nn.BatchNorm scales, all biases, and every bare nn.Parameter such as the mlp_plr periodic-embedding tensors) is excluded. Set to 0 to recover plain AdamW with no regularization.
cell_weight (float, default: 1.0) – Multipliers on the two reconstruction + commitment loss branches.
neighborhood_weight (float, default: 1.0) – Multipliers on the two reconstruction + commitment loss branches.
k_neighbors (int, default: 20) – Neighbors used to build the neighborhood feature.
neighborhood_aggregation (str, default: 'weighted_mean') – One of {"mean", "weighted_mean", "max", "gaussian", "inverse_square"}.
spatial_graph (str, default: 'knn_radius') – Spatial graph backend: {"knn", "knn_radius" (default), "radius", "delaunay", "alpha_complex", "gabriel", "rng"}.
radius (float | None, default: 50.0) – Radius in microns for the radius-capped graphs (spatial_graph="knn_radius" or "radius"). Default 50.0.
bandwidth (Optional[float], default: None) – Gaussian kernel bandwidth in microns for neighborhood_aggregation="gaussian".
lr_schedule (str, default: 'plateau') – "plateau" (default, released) or "warmup_cosine". The "warmup_cosine" schedule is applied PER OPTIMIZER STEP over total_steps = steps_per_epoch * num_epochs: a linear warmup over the first warmup_frac of the steps, then a cosine decay to min_lr at the final step (Goyal 2017; Nicheformer 2025 uses warmup + cosine at spatial scale).
warmup_frac (float, default: 0.03) – Fraction of total optimizer steps spent linearly ramping the LR from ~0 up to the peak (base) LR before the cosine decay begins, for "warmup_cosine". 0.03 (default). 0.0 reproduces a pure per-step cosine anneal with no warmup (back-compatible). Ignored by the plateau schedule.
warmup_steps (int, default: 0) – min_lr is the learning-rate floor the cosine decays to at the final step for "warmup_cosine" (0.0 decays fully to zero). warmup_steps is a legacy epoch-based field retained for API compatibility; the per-step "warmup_cosine" schedule uses warmup_frac for its warmup length.
min_lr (float, default: 0.0) – min_lr is the learning-rate floor the cosine decays to at the final step for "warmup_cosine" (0.0 decays fully to zero). warmup_steps is a legacy epoch-based field retained for API compatibility; the per-step "warmup_cosine" schedule uses warmup_frac for its warmup length.
decoupled_weight_decay (bool, default: True) – Use selective (decoupled) AdamW weight decay: decay ONLY the weight matrices of nn.Linear / nn.Conv* (and the attention QKV projection), and exclude everything else (both VQ codebook embeddings, the molecule-set gene embedding, all norm scales, all biases, and all bare nn.Parameter tensors including the mlp_plr freq / emb_w / emb_b periodic-embedding parameters). True (default). When False, a single AdamW group applies weight_decay uniformly to every parameter (including the codebooks), which is discouraged for a VQ tokenizer because decaying codebook vectors toward zero shrinks code norms and distorts nearest-code assignment geometry.
spatial_loss_type (str, default: 'none') – Optional spatial-coherence regularizer on the cell embedding ("laplacian" / "contrastive" / "codebook_consistency"), its weight (0 disables, the default), and its neighbor count. The graph is built within each mini-batch and restricted to same-sample pairs, so it is most effective with a large batch_size.
spatial_loss_weight (float, default: 0.0) – Optional spatial-coherence regularizer on the cell embedding ("laplacian" / "contrastive" / "codebook_consistency"), its weight (0 disables, the default), and its neighbor count. The graph is built within each mini-batch and restricted to same-sample pairs, so it is most effective with a large batch_size.
spatial_loss_k (int, default: 6) – Optional spatial-coherence regularizer on the cell embedding ("laplacian" / "contrastive" / "codebook_consistency"), its weight (0 disables, the default), and its neighbor count. The graph is built within each mini-batch and restricted to same-sample pairs, so it is most effective with a large batch_size.
normalize (bool, default: True) – Apply sc.pp.normalize_total / sc.pp.log1p to adata.X before training. Skipped if the AnnData is already log-normalized.
log1p (bool, default: True) – Apply sc.pp.normalize_total / sc.pp.log1p to adata.X before training. Skipped if the AnnData is already log-normalized.
val_fraction (float, default: 0.0) – Fraction of cells held out for a validation loss. 0 (default) trains on all cells and monitors the training loss (released behavior).
early_stopping_patience (Optional[int], default: None) – Stop if the monitored loss does not improve for this many epochs. None (default) disables early stopping.
grad_clip (Optional[float], default: None) – Max global gradient norm (torch.nn.utils.clip_grad_norm_). None (default) disables clipping.
amp (bool, default: False) – Enable CUDA automatic mixed precision. False (default) keeps full fp32 for bit-stable determinism.
resume_from (Optional[str], default: None) – Optional checkpoint .pt path to resume model weights from.
save_best (bool, default: True) – If True and a validation split is used, also write best_checkpoint.pt at the epoch with the lowest monitored loss.
seed (int, default: 9) – Seed passed to nicheverse.utils.seed_everything().
log_every (int, default: 10) – Emit a batch-level log line every log_every mini-batches.
deterministic (bool, default: True) – Forwarded to seed_everything().
num_workers (int, default: 0) – DataLoader workers.
ddp (bool, default: False) – Opt-in multi-GPU DistributedDataParallel. When True and the process was launched under torchrun (WORLD_SIZE > 1), the model is wrapped in DDP and the global batch is split across ranks (per-GPU batch = batch_size // world_size), so optimization matches a single-GPU run on the full global batch within numerical tolerance. The EMA VQ codebook statistics are all-reduced so the codebook is rank-correct. Default False and a no-op when WORLD_SIZE == 1 (single-GPU byte-identical). See DETERMINISM.md for the one documented residual (dead-code reset RNG).
compile_model (bool, default: False) – Opt-in torch.compile(model) for a faster (accuracy-preserving) forward / backward. Default False. The compiled graph is numerically equivalent to eager for these ops; only kernel fusion changes.
pin_memory (Optional[bool], default: None) – Explicit DataLoader pin_memory. None (default) auto-selects True on CUDA, preserving the released behavior. Verified accuracy neutral (byte-identical loss with num_workers=0).
persistent_workers (Optional[bool], default: None) – Explicit DataLoader persistent_workers. None (default) is False so the epoch-by-epoch shuffle order matches num_workers=0 exactly; set True to save worker respawn overhead at the cost of a different (still valid) shuffle for epochs after the first.
prefetch_factor (Optional[int], default: None) – DataLoader prefetch_factor (only applies when num_workers > 0). None (default) uses the PyTorch default (2). Verified accuracy neutral.
device_resident (bool, default: True) –
GPU-resident dataset that eliminates the per-batch host->device copy. Now default ON (automatic): the trainer uses the GPU-resident path whenever a CUDA GPU is present, DDP is off, and the dataset fits the GPU’s free-memory budget (see _device_resident_fits()); otherwise it transparently falls back to the CPU DataLoader. Set False to force the released CPU-DataLoader path. When active, the feature tensors (cell_features, neighborhood_features and, in the count modes, recon_target / niche_count_target) are moved to the GPU ONCE after dataset construction; the DataLoader then yields only the seeded batch INDEX tensor and the features are gathered on the GPU inside the loop (cb = cell_features_gpu[idx]). This is accuracy neutral: the SAME seeded shuffle/BatchSampler produces the SAME batch indices in the SAME order as the CPU path, so for a fixed seed the per-batch loss is equivalent to floating-point tolerance. Requires CUDA; on CPU it is a no-op (falls back to the CPU path). The fit-check is additive: the resident tensors plus a fixed ~8 GiB working-set margin must fit under device_resident_mem_fraction (default 0.9) of the memory budget, where the budget is the FREE GPU memory (torch.cuda.mem_get_info) when queryable, else the device total. If they do not fit (wide panels, e.g. CosMx 21731 genes or a big merged Xenium), a clear message is logged and the normal CPU DataLoader path is used instead (never OOMs silently). Because the resident tensors cannot be forked to DataLoader workers, num_workers is forced to 0 when this path is active (logged). Not currently combined with DDP: under DDP this flag is ignored (logged) and the CPU path is used.

Caveat: strictly byte-identical reproduction of a released checkpoint should pass device_resident=False. The on-GPU gather shifts floating-point accumulation by ~1e-5 (accuracy neutral, but not bit-identical to the CPU-DataLoader run that produced the release).
device_resident_mem_fraction (float, default: 0.9) – Fraction of the GPU memory budget the resident tensors + working margin may occupy in the device_resident fit-check. 0.9 (default). Lower it to be more conservative on a shared card.

Notes

num_workers=0 (the default) is byte-identical to the released training run. pin_memory and prefetch_factor are verified to leave the loss trajectory byte-identical. num_workers > 0 is a throughput knob for I/O-bound cohorts (large on-disk data); on small in-memory cohorts it can shift the trajectory to a different (equally valid) one, so it is not guaranteed bit-identical to num_workers=0. Prefer num_workers=0 when exact reproduction of the released run is required.

nicheverse.train_model(adata, checkpoint_dir, model_config=None, train_config=None, sample_col='sample_id', device=None)#

Train HierarchicalVQVAE on adata and persist checkpoint + outputs.

Parameters:

adata (AnnData) – Input AnnData. Required: adata.obsm['spatial'] (microns) and adata.obs[sample_col]. adata.X is expected to be raw counts; per-cell normalization and log1p are applied per train_config. adata is copied before mutation; the original is left untouched.
checkpoint_dir (str | Path) – Output directory for the checkpoint, embeddings, loss curve, environment snapshot, and annotated AnnData. Created if missing.
model_config (Optional[ModelConfig], default: None) – Optional ModelConfig. Defaults to a config with input_dim = adata.X.shape[1] and gene_names = adata.var_names.
train_config (Optional[TrainConfig], default: None) – Optional TrainConfig. Defaults to the constructor defaults.
sample_col (str, default: 'sample_id') – Column in adata.obs used to partition cells for the per-sample graph.
device (Optional[str], default: None) – Optional explicit device string (e.g. "cuda:0"). Defaults to the first CUDA device if available, else CPU.

Returns:

The trained model and the annotated AnnData copy with code indices and embeddings attached.

Return type:

(HierarchicalVQVAE, AnnData)

Raises:

ValueError – If adata.obsm['spatial'] or adata.obs[sample_col] is missing, or if model_config.input_dim disagrees with adata.X.shape[1].

nicheverse.mae_pretrain(adata, encoder_type='mlp', hidden=(256, 128), embedding_dim=64, mask_ratio=0.5, num_epochs=20, batch_size=2048, lr=0.001, normalize=True, log1p=True, seed=9, device=None)#

Pretrain an encoder by masked-gene reconstruction; return the pretrained encoder.

A fraction mask_ratio of genes is zeroed per cell; the encoder plus a light decoder reconstruct the full profile with the loss taken on masked positions only, biasing the encoder toward coexpression structure. The returned encoder has the same architecture as a model built with the matching encoder_type / hidden / embedding_dim and can initialize its cell_encoder via load_state_dict.

Parameters:

encoder_type (str, default: 'mlp')
embedding_dim (int, default: 64)
mask_ratio (float, default: 0.5)
num_epochs (int, default: 20)
batch_size (int, default: 2048)
lr (float, default: 0.001)
normalize (bool, default: True)
log1p (bool, default: True)
seed (int, default: 9)
device (Optional[str], default: None)

Prediction#

nicheverse.predict_codes(adata, checkpoint, output_path=None, sample_col='sample_id', k_neighbors=None, neighborhood_aggregation=None, spatial_graph=None, radius=<object object>, bandwidth=None, batch_size=2048, normalize=True, log1p=True, device=None, return_embeddings=True, config=None, seed=9, deterministic=True)#

Assign cell and neighborhood codebook indices to adata using a checkpoint.

Parameters:

adata (AnnData) – Input AnnData with raw counts (or already-normalized data, in which case set normalize=False and log1p=False). Must carry adata.obsm['spatial'] and adata.obs[sample_col].
checkpoint (str | Path | HierarchicalVQVAE) – Either a path to a .pt file written by nicheverse.train_model() or an already-loaded HierarchicalVQVAE.
output_path (Union[str, Path, None], default: None) – Optional .h5ad path to write the annotated AnnData. Parent directories are created if missing.
sample_col (str, default: 'sample_id') – Column in adata.obs used to partition cells for per-sample k-NN.
k_neighbors (Optional[int], default: None) – Neighbors per cell. None (default) inherits the training-time value from the checkpoint’s train_config.json so the neighborhood codebook behaves as intended; pass an int to override.
neighborhood_aggregation (Optional[str], default: None) – One of {"mean", "weighted_mean", "max", "gaussian", "inverse_square"}. None (default) inherits the training-time value from the checkpoint.
spatial_graph (Optional[str], default: None) – Spatial-graph backend. None (default) inherits the training-time value from the checkpoint (e.g. "knn_radius"); pass a string to override. Building the neighborhood graph with a different backend than training silently produces mismatched niche codes, so inheriting is the safe default.
radius (float | None, default: <object object at 0x7fa948ac2d70>) – Radius in microns for radius-based graphs. Left un-passed it inherits the training-time value from the checkpoint; pass an explicit value (including None for graphs that take no radius) to override.
bandwidth (Optional[float], default: None) – Gaussian kernel bandwidth in microns; must match training when neighborhood_aggregation="gaussian". None inherits the training-time value from the checkpoint.
batch_size (int, default: 2048) – Inference mini-batch size.
normalize (bool, default: True) – Apply sc.pp.normalize_total / sc.pp.log1p before inference. Set both to False if your AnnData is already log-normalized.
log1p (bool, default: True) – Apply sc.pp.normalize_total / sc.pp.log1p before inference. Set both to False if your AnnData is already log-normalized.
device (Optional[str], default: None) – Optional explicit device string. Defaults to CUDA if available.
return_embeddings (bool, default: True) – If True, write X_cell_embedding and X_neighborhood_embedding to adata.obsm.
config (Optional[ModelConfig], default: None) – Optional ModelConfig for legacy bare state dict checkpoints.
seed (int, default: 9) – Seed for the random ops within seed_everything().
deterministic (bool, default: True) – Forwarded to seed_everything().

Returns:

Annotated copy of the input with obs['cell_codebook_idx'] and obs['neighborhood_codebook_idx'] (plus embeddings in obsm if return_embeddings).

Return type:

AnnData

Raises:

ValueError – If required obs / obsm fields are missing, or if the gene panel does not match the checkpoint.

Losses#

Optional spatial-coherence and codebook regularizers added to the reconstruction objective.

Reconstruction likelihoods and optional spatial-coherence losses.

This module holds two families of loss functions used by nicheverse.train_model():

Count reconstruction likelihoods, nb_nll() and poisson_nll(), selected by ModelConfig.recon when the decoder models raw counts instead of Gaussian (MSE) residuals.
Spatial-coherence regularizers, laplacian_smoothness(), spatial_contrastive() and codebook_consistency(), collected in the SPATIAL_LOSSES registry. Each operates on a batch of latent vectors z of shape (B, D) and the matching cell coordinates coords of shape (B, 2) (microns), building an intra-batch k-nearest-neighbor graph on the fly. They are opt-in: training applies one only when TrainConfig.spatial_loss_weight > 0 (default 0), so the released training path is unaffected.

The spatial losses encourage the cell encoder to produce representations that vary smoothly across physical space, biasing neighboring cells toward the same (or nearby) discrete cell-state codes without hard-coding spatial information into the codebook itself.

Symbols#

B batch size, D latent dimensionality, k neighbors per cell.

nicheverse.losses.SPATIAL_LOSSES#

Registry mapping a spatial-loss name ("laplacian", "contrastive", "codebook_consistency") to its function. TrainConfig.spatial_loss_type selects the entry used during training.

Type:: dict[str, Callable]

nicheverse.losses.NICHE_SPATIAL_LOSSES = frozenset({'graph_tv'})#: Spatial-loss names applied to the NICHE (neighborhood) latent instead of the cell latent. "graph_tv" (L1 fused-lasso total variation) smooths niche assignments across the spatial graph, so the trainer feeds it the neighborhood encoder output; all other entries operate on the cell latent.

nicheverse.losses.gaussian_nll(target, pred)#

Gaussian reconstruction loss (mean squared error), the default cell/niche head.

Up to an additive constant and a factor, the negative log-likelihood of a Gaussian with fixed unit variance is the mean squared error between the prediction and the target, so this is the MSE reconstruction term named as a likelihood for uniform, by-name dispatch alongside nb_nll(), poisson_nll(), and dirichlet_multinomial_nll(). It is the released default for both the cell branch (on the log1p expression) and the niche branch (on the aggregated composition), and is numerically identical to torch.nn.functional.mse_loss(pred, target) (mean reduction).

Parameters:

target (torch.Tensor) – Reconstruction target (any shape).
pred (torch.Tensor) – Predicted reconstruction, same shape as target.

Returns:

Scalar mean squared error.

Return type:

nicheverse.losses.nb_nll(x, cr, log_theta, eps=1e-08, library=None)#

Negative-binomial negative log-likelihood (scVI mean-dispersion parameterization).

Overdispersed count reconstruction loss selected by ModelConfig.recon="nb". The decoder emits per-gene logits cr that are softmaxed into a proportion over the panel and scaled by the per-cell library size to give the NB mean mu; exp(log_theta) is the per-gene inverse-dispersion. This penalizes the model for assigning low probability to the observed integer counts and, unlike MSE, respects the mean-variance relationship of molecular counts, which better matches sparse Xenium panels. Used in place of the Gaussian (MSE) cell reconstruction term when raw counts are modeled directly.

Parameters:

x (torch.Tensor) – Observed raw counts, shape (B, N) (N genes). Not normalized.
cr (torch.Tensor) – Per-gene decoder logits, shape (B, N); softmaxed to proportions and scaled by library to form the NB mean mu.
log_theta (torch.Tensor) – Per-gene log inverse-dispersion, shape (N,) (a learned model.cell_log_theta parameter).
eps (float, default=1e-8) – Numerical floor added inside logs.
library (torch.Tensor, optional) – Per-cell multiplier for the softmax proportion (the library / GLM offset), shape (B,) or (B, 1). When None (default, byte-identical to the released path) the observed count sum x.sum(1) is used, which is the standard scVI library. Pass an explicit per-cell size factor to override it (the trainer passes the observed total count of the raw-count target so the softmax proportion is scaled to the correct count scale even though the encoder input was the log1p expression).

Returns:

Scalar mean NB negative log-likelihood over the batch.

Return type:

nicheverse.losses.poisson_nll(x, cr, eps=1e-08, library=None)#

Poisson negative log-likelihood with a library-size-scaled softmax mean.

Count reconstruction loss selected by ModelConfig.recon="poisson". Like nb_nll() the decoder logits cr are turned into a proportion and scaled by the per-cell library size to form the Poisson rate mu, but with no separate dispersion parameter (variance equals mean). It is the simpler, equidispersed alternative for count reconstruction; prefer nb_nll() when the data are overdispersed.

Parameters:

x (torch.Tensor) – Observed raw counts, shape (B, N).
cr (torch.Tensor) – Per-gene decoder logits, shape (B, N); softmaxed and scaled by the library size to form the rate mu.
eps (float, default=1e-8) – Numerical floor inside the log.
library (torch.Tensor, optional) – Per-cell multiplier for the softmax proportion (the library / GLM offset), shape (B,) or (B, 1). When None (default) the observed count sum x.sum(1) is used; pass an explicit per-cell size factor to override it. See nb_nll().

Returns:

Scalar mean Poisson negative log-likelihood over the batch.

Return type:

nicheverse.losses.dirichlet_multinomial_nll(target, logits, log_alpha, eps=1e-08)#

Dirichlet-multinomial negative log-likelihood for a compositional target.

A proper likelihood for a niche composition vector: the decoder logits are softmaxed into a mean composition p and scaled by a learned per-feature concentration exp(log_alpha) to form the Dirichlet parameters alpha = p * sum(exp(log_alpha)) (a mean/precision parameterization). The target is treated as a (possibly fractional) count vector c with total N = c.sum(1), and the Dirichlet-multinomial log-density is evaluated with lgamma (well-defined for non-integer c, so it handles the continuous aggregated-neighbor composition without rounding). Unlike an MSE on the composition, it respects that the target is a normalized, overdispersed proportion. Opt-in via ModelConfig.niche_recon="dirichlet_multinomial".

Parameters:

target (torch.Tensor) – COUNT-SCALE compositional target, shape (B, N), non-negative. Its per-row sum is used as the multinomial total N, so this must be on the raw-count scale (e.g. the weighted-mean aggregation of the raw neighbor counts), NOT a log1p mean whose row sum has no count interpretation. Fractional values are allowed (the density is evaluated with lgamma), so a continuous count-scale aggregate is fine without rounding.
logits (torch.Tensor) – Per-feature decoder logits, shape (B, N); softmaxed to the mean composition.
log_alpha (torch.Tensor) – Per-feature log-concentration, shape (N,) (a learned model.niche_log_alpha parameter). sum(exp(log_alpha)) sets the overall precision.
eps (float, default=1e-8) – Numerical floor.

Returns:

Scalar mean Dirichlet-multinomial NLL over the batch.

Return type:

nicheverse.losses.bernoulli_detection_bce(x, logits)#

Binary-cross-entropy detection hurdle on the 0-vs-nonzero mask of counts.

A per-gene Bernoulli likelihood over whether each gene is DETECTED (count > 0) in a cell, using logits as the detection logits. Summing this with a count NLL forms a hurdle model that separates the probability of detection from the conditional count magnitude, which better fits the excess zeros of sparse imaging panels (Xenium). Opt-in via ModelConfig.detection_weight > 0.

Parameters:

x (torch.Tensor) – Observed raw counts, shape (B, N). The target mask is (x > 0).
logits (torch.Tensor) – Per-gene detection logits, shape (B, N) (the cell decoder output is reused as the logits, so no extra head is needed).

Returns:

Scalar mean BCE (summed over genes, averaged over cells).

Return type:

nicheverse.losses.graph_total_variation(z, coords, k=6)#

Graph total-variation (L1 fused-lasso) penalty on latents across neighbors.

The L1 analogue of laplacian_smoothness(): it penalizes the mean L1 (not squared-L2) difference between each cell’s latent and its k spatial nearest neighbors, mean_i mean_j |z_i - z_j|_1. The L1 form is the fused lasso / anisotropic total variation, which promotes piecewise-constant latent fields (sharp niche boundaries with flat interiors) rather than the smooth gradients an L2 penalty prefers. Registered as spatial loss "graph_tv" and applied to the niche latent over the per-sample in-batch graph; it is opt-in (TrainConfig.spatial_loss_weight > 0). Note that random mini-batches contain few true spatial edges, so a spatial-contiguous sampler would strengthen this term (not built here).

Parameters:

z (torch.Tensor) – Latent vectors, shape (B, D).
coords (torch.Tensor) – Cell coordinates in microns, shape (B, 2).
k (int, default=6) – Nearest neighbors per cell (excluding itself). Capped at B - 1.

Returns:

Scalar mean L1 total-variation penalty over neighbor pairs.

Return type:

nicheverse.losses.codebook_consistency(z, coords, k=6, margin=0.5)#

Margin hinge penalty on latent distance between spatial neighbors.

Computes the Euclidean distance between each cell’s latent and each of its k spatial neighbors and penalizes only the amount by which that distance exceeds margin (a hinge / ReLU). This is a softer, tolerance-banded version of laplacian_smoothness(): neighbors are allowed to differ up to margin free of charge, so genuine local heterogeneity (e.g. a tumor boundary) is not over-smoothed, while grossly divergent neighbors are pulled together to stabilize the codebook assignment of adjacent cells.

Parameters:

z (torch.Tensor) – Latent vectors, shape (B, D).
coords (torch.Tensor) – Cell coordinates in microns, shape (B, 2).
k (int, default=6) – Spatial neighbors per cell. Valid range >= 1.
margin (float, default=0.5) – Distance tolerance below which neighbor pairs incur no penalty. Larger values permit more local variation. Must be >= 0.

Returns:

Scalar mean hinge penalty over neighbor pairs.

Return type:

nicheverse.losses.laplacian_smoothness(z, coords, k=6)#

Graph-Laplacian smoothness penalty on latents across spatial neighbors.

Penalizes the mean squared Euclidean difference between each cell’s latent and those of its k spatial nearest neighbors (STAGATE-style Laplacian regularization). Minimizing it drives spatially adjacent cells toward similar embeddings, which in turn makes them more likely to quantize to the same cell state code. It is unbounded below only trivially (constant z gives 0), so it is used with a small weight alongside the reconstruction objective.

Parameters:

z (torch.Tensor) – Latent vectors, shape (B, D).
coords (torch.Tensor) – Cell coordinates in microns, shape (B, 2), used to build the intra-batch neighbor graph.
k (int, default=6) – Number of nearest neighbors per cell (excluding itself). Effectively capped at B - 1. Valid range >= 1.

Returns:

Scalar smoothness penalty (mean over cells and neighbors).

Return type:

nicheverse.losses.spatial_contrastive(z, coords, k=6, temperature=0.07, num_negatives=64)#

InfoNCE contrastive loss pulling spatial neighbors together, pushing negatives apart.

L2-normalizes the latents, then for each cell forms a positive score as the mean cosine similarity to its k spatial neighbors and negative scores against num_negatives randomly sampled in-batch cells, all scaled by 1 / temperature, and returns the mean InfoNCE loss. Unlike laplacian_smoothness() it is collapse-resistant: the negatives prevent the encoder from mapping everything to one point, so it shapes a spatially organized yet still discriminative latent space for the cell quantizer.

Parameters:

z (torch.Tensor) – Latent vectors, shape (B, D).
coords (torch.Tensor) – Cell coordinates in microns, shape (B, 2).
k (int, default=6) – Spatial neighbors treated as positives. Valid range >= 1.
temperature (float, default=0.07) – Softmax temperature on the similarity logits; lower sharpens the contrast. Must be > 0.
num_negatives (int, default=64) – Number of random in-batch negatives sampled per cell (sampled with replacement, so a positive may occasionally recur as a negative).

Returns:

Scalar InfoNCE loss averaged over cells.

Return type:

Annotation#

Turn learned codes into labeled cell states and spatial niches, grounded in per-code marker and DEG evidence and, optionally, primary literature. The end-to-end harness is annotate_codebook(), which builds the evidence, proposes labels, gates them against markers and lab rules, scores calibration, and returns an AnnotationResult. The individual stages (code_evidence(), annotate_codes(), the verify and evaluate helpers) are also callable on their own.

Harness#

nicheverse.annotate.annotate_codebook(adata, code_col, *, config=None, kind='cell', celltype_col=None, groundtruth_col=None, out_dir=None, run_meta=None, resolver=None, lineage_map=None)#

Run the full labeler -> gate -> refuter -> reconcile -> review -> provenance pipeline.

Parameters:

adata – AnnData and the obs column holding the code index (cell or niche codes).
code_col (str) – AnnData and the obs column holding the code index (cell or niche codes).
config (Optional[AnnotationConfig], default: None) – AnnotationConfig; set config.context to a ProjectContext to ground labels and supply the allowed vocabulary. config.refuter toggles the adversarial pass; config.min_marker_precision and config.confidence_review_threshold tune the gate and the review split.
kind (str, default: 'cell') – "cell" (default) or "niche". For niches, celltype_col (cell-level labels) is required for the composition-primary evidence.
celltype_col (Optional[str], default: None) – Cell-type label column used to build niche composition (kind='niche' only).
groundtruth_col (Optional[str], default: None) – Optional independent label column; when given, each code is scored against its majority ground-truth label and the run is calibrated.
out_dir (Optional[str], default: None) – When given, the provenance manifest + scorecard CSV are written here.
run_meta (Optional[dict], default: None) – Run-level metadata recorded verbatim in the manifest (model, seed, timestamp, …).
resolver (default: None) – Injectable citation resolver (id) -> record|None handed to verify.gate() (keeps citation checks offline in tests).
lineage_map (Optional[dict], default: None) – Optional label -> lineage (or (lineage, depth)) map for evaluate.score_code().

Return type:

AnnotationResult

class nicheverse.annotate.AnnotationResult(labels_df, review_df, scorecard_df=None, calibration=None, manifest_path=None, records=<factory>, kind='cell')#

Bases: object

Everything the harness produces for one codebook annotation run.

labels_df#: One row per code: final_label, compartment, confidence (post-penalty), passed (gate), refuter_agree, and flags.

review_df#: Subset of codes that did not auto-accept (gate failed, refuter disagreed, or low confidence), with a review_reason column, for manual sign-off.

scorecard_df#: Per-code grading table (from evaluate.score_code()) when a ground-truth column was supplied, else None.

calibration#: Confidence-vs-correctness calibration dict when scored, else None.

manifest_path#: Path to the provenance manifest JSON when out_dir was given, else None.

records#: The per-code provenance records (kept in memory for inspection / re-use).

kind#: "cell" or "niche".

Parameters:

labels_df (DataFrame)
review_df (DataFrame)
scorecard_df (Optional[DataFrame], default: None)
calibration (Optional[dict], default: None)
manifest_path (Optional[str], default: None)
records (list, default: <factory>)
kind (str, default: 'cell')

label_map()#

code -> final_label mapping for attaching onto obs.

Return type:: dict

attach(adata, key_added='celltype_annot')#

Map final labels onto adata.obs[key_added] via annotate.attach_labels().

The code column is inferred from the label DataFrame’s index name when present.

Parameters:: key_added (str, default: 'celltype_annot')

class nicheverse.annotate.AnnotationConfig(provider='anthropic', model=None, tissue='', with_literature=False, context_cols=(), marker_context='cell type marker', top_markers=25, top_degs=20, refine=True, cluster_context=False, api_key=None, context=None, refuter=True, n_candidates=3, min_marker_precision=0.5, confidence_review_threshold=0.6)#

Bases: object

Settings for annotate_codes().

Parameters:

provider (str, default: 'anthropic') – LLM backend (anthropic / openai / ollama) and model id.
model (Optional[str], default: None) – LLM backend (anthropic / openai / ollama) and model id.
tissue (str, default: '') – Free-text tissue / disease context (e.g. “human clear cell RCC, brain metastasis”).
with_literature (bool, default: False) – If True, search PubMed for the top markers of each code and include the hits.
context_cols (tuple[str, ...], default: ()) – obs columns to summarize per code (e.g. ("site_class", "sample_id")).
refine (bool, default: True) – Run a second pass that reconciles labels across codes (resolves duplicates / conflicts).
cluster_context (bool, default: False) – Hierarchically cluster codes first and give the reconciliation pass that structure, so similar codes get consistent compartments and distinguishing qualifiers.
marker_context (str, default: 'cell type marker')
top_markers (int, default: 25)
top_degs (int, default: 20)
api_key (Optional[str], default: None)
context (Optional[object], default: None)
refuter (bool, default: True)
n_candidates (int, default: 3)
min_marker_precision (float, default: 0.5)
confidence_review_threshold (float, default: 0.6)

Project context and priors#

class nicheverse.annotate.ProjectContext(name='', species='', disease='', tissue='', platform='', panel='', sites=<factory>, expected_cell_types=<factory>, expected_niches=<factory>, site_col='', patient_col='', context_cols=<factory>, references=<factory>, notes='')#

Bases: object

Structured, user-supplied context grounding LLM codebook annotation.

All fields are optional. platform is validated on construction: a known non-IMST platform raises ValueError (nicheverse is imaging-based only), an unknown non-empty platform warns but is kept, and an empty platform is allowed.

Parameters:

name (str, default: '')
species (str, default: '')
disease (str, default: '')
tissue (str, default: '')
platform (str, default: '')
panel (str, default: '')
sites (list[str], default: <factory>)
expected_cell_types (list[CellTypePrior], default: <factory>)
expected_niches (list[NichePrior], default: <factory>)
site_col (str, default: '')
patient_col (str, default: '')
context_cols (list[str], default: <factory>)
references (list[str], default: <factory>)
notes (str, default: '')

classmethod from_dict(d)#

Build from a plain dict, tolerating unknown keys and coercing nested priors.

Robust to messy YAML: None -> an empty context; unknown top-level keys are dropped; expected_cell_types / expected_niches may be given as dataclass instances, plain dicts, or any mix; a malformed nested entry (not a dict and not the matching dataclass, e.g. a bare string) is skipped rather than raising.

Parameters:: d (dict | None)
Return type:: ProjectContext

classmethod from_yaml(path)#

Load from a YAML file via yaml.safe_load(), then from_dict().

Parameters:: path (str)
Return type:: ProjectContext

to_dict()#

Return a plain, YAML-serializable dict (nested priors expanded).

Return type:: dict

to_yaml(path)#

Write the context to a YAML file.

Parameters:: path (str)
Return type:: None

to_prompt_block(max_cell_types=12, max_niches=8, max_markers=8)#

Render a compact, deterministic plain-text context block for an LLM prompt.

Omits empty fields, truncates long lists gracefully, and stays well under ~1500 chars for modest inputs. Uses label: value lines and numbered items (no markdown - bullets). The rendered block is guaranteed free of em-dashes and of leading markdown bullet markers (-, *, + followed by a space) even when user-supplied field values (names, markers, notes, references) contain them: _clean() rewrites those characters so the block always complies with the repo’s prose rules. Stored field values are never mutated (only the rendered copy is), so (de)serialization roundtrips are unaffected.

Parameters:

max_cell_types (int, default: 12)
max_niches (int, default: 8)
max_markers (int, default: 8)

Return type:

str

static write_template(path)#

Write an annotated example project_context.yaml the user can fill in.

Copies the shipped template next to this module when present; otherwise writes the same content inline (so it does not hard-depend on package-data config).

Parameters:: path (str)
Return type:: None

class nicheverse.annotate.CellTypePrior(name, markers=<factory>, notes='')#

Bases: object

A user-provided expectation for a cell type present in the study.

Parameters:

name (str)
markers (list[str], default: <factory>)
notes (str, default: '')

class nicheverse.annotate.NichePrior(name, description='', expected_cell_types=<factory>)#

Bases: object

A user-provided expectation for a multicellular spatial niche.

Parameters:

name (str)
description (str, default: '')
expected_cell_types (list[str], default: <factory>)

Evidence and labeling#

nicheverse.annotate_codes(adata, code_col, config=None, **kwargs)#

Annotate every code in adata.obs[code_col] with an LLM, grounded in evidence.

Returns a DataFrame indexed by code with label / compartment / confidence / rationale / key_markers / citations (and label_refined when config.refine).

Parameters:

adata (AnnData)
code_col (str)
config (Optional[AnnotationConfig], default: None)

Return type:

DataFrame

nicheverse.code_evidence(adata, code_col, *, extra_cols=(), top_markers=30, top_degs=30, layer=None)#

Assemble per-code evidence for annotation.

Parameters:

adata (AnnData) – AnnData carrying the code assignment in obs[code_col].
code_col (str) – Column of obs with the code index (e.g. cell_codebook_idx).
extra_cols (tuple[str, ...], default: ()) – obs columns to summarize the per-code distribution over (e.g. ("site_class", "sample_id")).
top_markers (int, default: 30) – Number of top markers per code (ranked by z-score across codes).
top_degs (int, default: 30) – Number of one-vs-rest DEGs per code (rank_genes_groups).
layer (Optional[str], default: None) – Optional layers key to read expression from; defaults to X.

Returns:

{code: {code, n_cells, frac, top_markers, top_degs, dist_<col>}}. top_markers is a list of (gene, z); top_degs a list of (gene, log2fc, padj).

Return type:

dict

nicheverse.annotate.annotate_niches(adata, niche_col, celltype_col, config=None, **kwargs)#

Annotate spatial-niche codes by their cell-type composition with an LLM.

Requires celltype_col (cell-level labels, e.g. from annotate_codes()). Returns a DataFrame indexed by niche code with label / dominant_types / confidence / rationale / composition.

Parameters:

adata (AnnData)
niche_col (str)
celltype_col (str)
config (Optional[AnnotationConfig], default: None)

Return type:

DataFrame

nicheverse.annotate.niche_evidence(adata, niche_col, celltype_col, *, extra_cols=(), top_markers=20, top_compositions=8)#

Per-niche evidence: the cell-type composition and enriched genes of each spatial niche.

A niche (neighborhood code) is characterized by the community of cell types that co-occur in it, so annotation needs celltype_col (typically the labels produced by annotating the cell codes) alongside marker enrichment.

Returns {niche: {code, n_cells, frac, composition [(cell_type, fraction)], top_markers [(gene, z)], dist_<col>}}.

Parameters:

adata (AnnData)
niche_col (str)
celltype_col (str)
extra_cols (tuple[str, ...], default: ())
top_markers (int, default: 20)
top_compositions (int, default: 8)

Return type:

dict[str, dict]

nicheverse.annotate.cluster_codes(adata, code_col, *, n_clusters=None, layer=None)#

Hierarchically group codes by mean-expression correlation.

Clusters similar codes together (correlation distance + average linkage) so they can be reviewed and annotated coarse-to-fine, as in the manual codebook pipeline. n_clusters defaults to about one cluster per eight codes.

Returns a DataFrame indexed by code with an integer cluster column.

Parameters:

adata (AnnData)
code_col (str)
n_clusters (Optional[int], default: None)
layer (Optional[str], default: None)

Return type:

DataFrame

nicheverse.annotate.code_context(adata, code_col, *, patient_col=None, context_cols=())#

Per-code patient / specimen context.

One row per code with n_cells, frac, the number of distinct patients (if patient_col given), and for each column in context_cols the dominant value and its fraction. This is the patient/specimen-context step of the codebook-review pipeline. Robust to absent columns (silently skipped).

Returns a DataFrame indexed by code with columns n_cells, frac, optional n_patients, and per context col <col>_dominant + <col>_dominant_frac.

Parameters:

adata (AnnData)
code_col (str)
patient_col (Optional[str], default: None)
context_cols (tuple[str, ...], default: ())

Return type:

DataFrame

nicheverse.annotate.code_groundtruth_concordance(adata, code_col, groundtruth_col)#

Cross-check codes against an independent ground-truth clustering/annotation.

Compares the learned codes in obs[code_col] to a separately-derived label column (e.g. a Leiden clustering or a manual annotation) in obs[groundtruth_col]. This is the Leiden cross-check step of the codebook-review pipeline.

Parameters:

adata (AnnData) – AnnData carrying both columns in obs.
code_col (str) – Column with the learned code index.
groundtruth_col (str) – Column with the independent ground-truth label.

Returns:

(crosstab, majority). crosstab is a DataFrame of code (rows) x ground-truth group (columns) normalized so each code row sums to ~1 (rows for codes with no valid ground-truth cells are all zero). majority has one row per code with columns code, majority_group, majority_frac, n_cells.

Return type:

tuple

nicheverse.annotate.write_evidence_bundle(adata, code_col, out_dir, *, groundtruth_col=None, patient_col=None, context_cols=(), top_markers=30, top_degs=30, layer=None)#

Compute the full per-code evidence bundle and dump it as CSVs.

Mirrors the documented 8-part codebook-review pipeline: per-code mean expression, z-score across codes, top markers, one-vs-rest DEGs, patient/specimen context, hierarchical code clustering, and (when a ground-truth column is given) the ground-truth crosstab + per-code majority.

Parameters:

adata (AnnData) – AnnData carrying the code assignment in obs[code_col].
code_col (str) – Column of obs with the code index.
out_dir (str) – Directory to write CSVs into (created if absent).
groundtruth_col (Optional[str], default: None) – Optional independent label column for the Leiden/ground-truth cross-check.
patient_col (Optional[str], default: None) – Optional column counted for n_patients in the context table.
context_cols (tuple[str, ...], default: ()) – obs columns to summarize (dominant value) in the context table.
top_markers (int, default: 30) – Numbers passed through to code_evidence().
top_degs (int, default: 30) – Numbers passed through to code_evidence().
layer (Optional[str], default: None) – Optional layers key for expression; defaults to X.

Returns:

Mapping of artifact name -> written CSV path.

Return type:

dict

nicheverse.annotate.attach_labels(adata, code_col, labels, key_added='celltype_annot')#

Map a code -> label mapping (dict or annotate_codes DataFrame) onto obs[key_added].

Parameters:

adata (AnnData)
code_col (str)
key_added (str, default: 'celltype_annot')

Return type:

AnnData

nicheverse.annotate.code_dotplot(adata, code_col, *, top_n=5, genes=None, save_path=None, layer=None, font_size=11)#

Plot each code’s top markers as a dot grid and save a vector PDF.

Parameters:

adata (AnnData) – AnnData and the obs column holding the code index.
code_col (str) – AnnData and the obs column holding the code index.
top_n (int, default: 5) – Top markers per code (by z-score) to include when genes is not given.
genes (Optional[list[str]], default: None) – Explicit gene list to show; otherwise the union of each code’s top markers.
save_path (Union[str, Path, None], default: None) – Output .pdf path. If omitted, the Matplotlib figure is returned.
layer (Optional[str], default: None)
font_size (int, default: 11)

Verification and scoring#

The verify stage checks proposed labels against the evidence and lab rules; the evaluate stage scores per-code agreement and calibrates confidence.

nicheverse.annotate.gate(label_dict, code_evidence, project_context, *, kind='cell', resolver=None, adata=None, code=None, z_thresh=1.0, min_marker_precision=0.5)#

Run every verification gate and synthesize an accept / penalize / reject decision.

POLICY. Exactly three failing CATEGORIES are evaluated, each contributing one flag and a fixed 0.34 to confidence_penalty:

MARKER. Markers were cited (n_cited > 0) AND marker['precision'] < min_marker_precision – the LLM cited genes not enriched in this code (hallucination or segmentation leakage). Citing nothing is not itself a marker failure.
CITATION. At least one citation has resolved is False – a parseable id the INJECTED resolver could not retrieve, i.e. a likely fabricated identifier. An unresolved offline / unknown id (resolved is None) does NOT fail the gate, so the default (no resolver, offline) path never fails a citation merely for being unresolved. An unparsed citation also does not fail here.
RULE. apply_lab_rules recorded any violation (a site-restricted mislabel or an unsupported composite).

INVARIANTS. confidence_penalty = min(1.0, 0.34 * n_failing_categories) and passed is (confidence_penalty == 0.0) exactly; every failing category contributes exactly one entry to flags (rule violations are appended verbatim), so an empty flags list is equivalent to passed is True. The penalty caps at 1.0 so a caller can downweight rather than hard-drop borderline calls.

adjusted_label is the lab-rule-adjusted label when a rule fired, else the label’s vocabulary-closest when it is out of vocab and a plausible near-match exists, else None (leave the label as-is).

A marker complaint about a NOVEL (out-of-vocab) label still fails the gate: the evidence must back whatever the LLM cited regardless of whether the label is in the vocabulary.

Returns:

{"passed": bool, "flags": [str,...], "marker": <marker_presence>,
 "citations": <check_citations>, "rules": <apply_lab_rules>,
 "vocab": <validate_vocabulary>, "adjusted_label": str|None,
 "confidence_penalty": float}

Parameters:

label_dict (dict)
code_evidence (dict)
kind (str, default: 'cell')
z_thresh (float, default: 1.0)
min_marker_precision (float, default: 0.5)

Return type:

nicheverse.annotate.marker_presence(cited_markers, code_evidence, *, z_thresh=1.0)#

Check which cited markers are actually enriched in the code’s evidence.

A cited marker is SUPPORTED if it appears in code_evidence['top_markers'] with z >= z_thresh OR as a positive DEG in code_evidence['top_degs']. Markers that clear neither bar are absent (hallucinated or segmentation leakage). Gene matching is case-insensitive.

Returns {"present": [...], "absent": [...], "precision": float, "n_cited": int}, where precision = len(present) / n_cited (and is 1.0 when nothing was cited, so an empty citation list is not itself penalized here).

Parameters:

code_evidence (dict)
z_thresh (float, default: 1.0)

Return type:

nicheverse.annotate.check_citations(citations, *, resolver=None, markers=None)#

Parse and (optionally) resolve each citation, flagging fabricated identifiers.

Each citation is parsed (see _parse_citation()) for a DOI or a PMID. resolver is a callable (identifier_or_query) -> record | None; when None it lazily falls back to PubMed via nicheverse.annotate.literature, but only if a network path is available. The resolver call is wrapped, so a resolver that raises (or an offline default with no network) yields a None record rather than propagating – this function never raises. Tests inject a dict-backed fake resolver and never touch the network.

Returns one dict per citation:

{"citation": raw, "id": parsed_or_None, "kind": "pmid"|"doi"|"unparsed",
 "resolved": bool|None, "supports": bool|None, "record": {...}|None}

resolved semantics (the key anti-hallucination distinction): :rtype: list[dict]

None – no verdict: the id was unparsed, OR the DEFAULT (offline) resolver ran and could not reach the network. An unknown id is NEVER accused of being fabricated.
True – an INJECTED or online resolver returned a record for the id.
False – an INJECTED resolver ran and returned nothing for a parseable id: the id is a likely fabrication. Only the injected path can yield False.

supports is True iff a retrieved record’s text mentions one of markers as a whole word, and None whenever no record was retrieved or markers is empty.

Return type:: list[dict]

nicheverse.annotate.apply_lab_rules(label_dict, code_evidence, project_context, *, adata=None, code=None)#

Encode the lab’s codebook conventions (context.ANNOTATION_RULES) as executable checks.

(a) Site-aware reassignment (rule 4): if the proposed label names a site-restricted (brain / CNS-only) type (matched as a whole word) but the code’s site distribution is not dominated by a permissive site – specifically the top-fraction site is non-permissive, OR a non-permissive site holds a material fraction (>= _SITE_MATERIAL_FRAC = 0.25) of the code’s cells – record a violation and propose the closest general label. A permissive site holding a small non-permissive tail (plausible segmentation leakage) is NOT flagged. When the site distribution is unavailable the check is skipped and a note is recorded.

(b) Composite label X/Y (rule 5): allowed only if BOTH components have marker support clearing a relative bar (rel >= 1.3x cohort mean) and an absolute bar (abs >= 0.05) in the evidence. Absolute / relative-to-cohort expression is NOT recoverable from the per-code evidence dict (which carries only across-code z-scores and one-vs-rest DEGs, not the cohort-mean expression the exact bars need), so the check APPROXIMATES that support by requiring each component’s canonical markers (or, absent priors, the component name itself) to be present via marker_presence(); this approximation is recorded in notes. A composite with fewer than two supported components collapses to its dominant component (most present markers, then higher precision). A site reassignment from (a) takes precedence over the collapsed label.

Returns {"violations": [...], "adjusted_label": str|None, "notes": [...]}. adjusted_label is None when no rule fired (i.e. the label is left as-is).

Parameters:

label_dict (dict)
code_evidence (dict)

Return type:

nicheverse.annotate.validate_vocabulary(label, project_context, *, kind='cell', allow_novel=True)#

Check a label against the project’s expected vocabulary with a novelty escape hatch.

For kind='cell' the allowed set is project_context.expected_cell_types names; for kind='niche' it is project_context.expected_niches names (any other kind is treated as 'cell'). in_vocab is set by an EXACT (case- and whitespace-insensitive) name match ONLY, so fuzzy matching can never manufacture a false in-vocab hit.

closest surfaces the nearest allowed label (see _closest()), gated by allow_novel:

allow_novel=True (default): only surface closest when it is a plausible near-match (similarity >= 0.6); an out-of-vocab label with no near neighbor returns closest=None and is left untouched by the caller (novel / uncertain codes are allowed, not errors).
allow_novel=False: novel labels are not tolerated, so closest is ALWAYS the nearest allowed label (when any vocabulary exists), letting the caller coerce an out-of-vocab label onto the closest expected type.

An exact in-vocab match always returns itself as closest. Returns {"in_vocab": bool, "closest": str|None, "kind": kind}.

Parameters:

label (str)
kind (str, default: 'cell')
allow_novel (bool, default: True)

Return type:

nicheverse.annotate.score_code(proposed, reference_label, code_evidence, *, lineage_map=None, gate_result=None, z_thresh=1.0, k=15)#

Score one proposed code label against an independent reference label.

Parameters:

proposed (dict) – The LLM’s proposed-label dict, e.g. {"label", "compartment", "confidence", "key_markers", "citations", ...}. key_markers may be a list or a ; / ,-joined string; missing keys are tolerated.
reference_label (str) – Independent ground-truth label for the code (e.g. the majority Leiden group from code_groundtruth_concordance()). May be None/empty (then label_match is None).
code_evidence (dict) – Per-code evidence dict from nicheverse.annotate.code_evidence() (top_markers=[(gene, z)], optional top_degs, dist_<col> keys).
lineage_map (Optional[dict], default: None) – Optional label -> broad lineage (or label -> (lineage, depth)) map that approximates the Cell Ontology tiers. See the module docstring.
gate_result (Optional[dict], default: None) – Optional dict from a decoupled gate (verify.gate), carrying marker precision, citation resolution, and rule violations. Passed in so this module never imports verify. Recognized keys (all optional): rules_pass (bool) or rule_violations (list); n_absent_markers / absent_markers; n_unresolved_citations / unresolved_citations; marker_precision.
z_thresh (float, default: 1.0) – A top_markers gene counts as enriched when its z-score >= this.
k (int, default: 15) – Marker precision/recall are computed against the code’s top-k markers.

Returns:

Flat scorecard with the dimensions (A) label_match (1.0/0.5/0.0/None), (B) granularity (“over_call”/”under_call”/”match”/None), (C) marker_precision / marker_recall, (D) rule_compliance (“pass”/”fail”), (E) hallucination counts n_absent_markers / n_unresolved_citations, plus the raw proposed label, reference, compartment, confidence, code, and any platform carried on the proposed dict.

Return type:

dict

nicheverse.annotate.calibration(scorecards, *, n_bins=4)#

Does the LLM’s stated confidence predict whether it was right?

Bins scorecards by confidence and reports the mean label_match per bin, plus the Spearman rank correlation between confidence and label_match. Records with a missing confidence or a None label_match are dropped from the correlation (and binning) but counted in n_total.

Returns {"spearman": float|None, "bins": {bin: {...}}, "n": int, "n_total": int}. A positive Spearman means higher confidence tracks higher correctness.

Parameters:

scorecards (list[dict])
n_bins (int, default: 4)

Return type: