Expand description
Retrieval-based imputation: pool a sparse reference’s full-feature counts onto query cells matched in a shared latent space.
The caller owns the latent semantics — both matrices must already be mapped onto whatever space makes their rows comparable under L2 (softmaxed topic proportions, L2-normalized embeddings, …). This module only does the model-agnostic part:
- Build an approximate kNN index over the reference rows.
- For each query row, find the K nearest reference rows and convert their L2 distances to softmax weights with temperature τ.
- Stream the reference backend’s columns chunk by chunk and
accumulate each query row’s weighted average of its neighbours’
counts into a dense
[N_query, G_ref]matrix.
An all-zero row means “no evidence” on BOTH sides: a zero query row is skipped (its imputed row stays zero), and a zero reference row is excluded from the index — left in, it would sit at one fixed coordinate capturing every otherwise-distant query.
The information-theoretic point of retrieval over a decoder readout: the imputed expression inherits the reference’s full-rank gene–gene covariance through the retrieved cells, rather than being confined to the rank of the model’s dictionary.
Structs§
Functions§
- retrieval_
impute - Weighted-kNN imputation of
ref_data’s features onto the query rows.