loki2.retrieve_from_sc

Retrieval from single-cell reference data for Loki2.

This module provides functionality for retrieving morphological embeddings from single-cell transcriptomic reference data using contrastive learning.

Module Contents

loki2.retrieve_from_sc.NormalizeT
loki2.retrieve_from_sc.NumericReductionT
class loki2.retrieve_from_sc.RetrievalResult

Result of a retrieval operation containing scores and indices.

Parameters:

  • scores – Similarity scores tensor of shape (M, K) where M is the number of queries and K is the number of neighbors.

  • indices – Indices tensor of shape (M, K) pointing to neighbors in the reference pool.

scores: torch.Tensor
indices: torch.Tensor
to_tuple() → Tuple[torch.Tensor, torch.Tensor]
save(path: str | pathlib.Path) → None

Persist scores/indices to disk via torch.save.

Parameters:

path – Path where to save the retrieval result.

classmethod load(path: str | pathlib.Path) → RetrievalResult

Restore a RetrievalResult saved by save.

Parameters:

path – Path to the saved retrieval result file.

Returns:

Loaded retrieval result.

Return type:

RetrievalResult

majority_vote(source_labels: Sequence | numpy.ndarray | torch.Tensor, *, weighted: bool | None = False, scores: torch.Tensor | None = None, temperature: float | None = None, return_counts: bool = False) → Tuple[numpy.ndarray, torch.Tensor | None]

Run k-NN majority voting using stored indices and optional scores.

Parameters:

  • source_labels – Labels for the reference pool.

  • weighted – Whether to weight votes by similarity scores. Defaults to False.

  • scores – Optional custom scores to use for weighting. Defaults to None.

  • temperature – Optional temperature for softmax weighting. Defaults to None.

  • return_counts – Whether to return vote counts. Defaults to False.

Returns:

  • Predicted labels

  • Vote counts if return_counts=True, otherwise None

Return type:

Tuple[np.ndarray, Optional[torch.Tensor]]

numeric_pool(source_values: Sequence | numpy.ndarray | torch.Tensor, *, weighted: bool | None = False, scores: torch.Tensor | None = None, temperature: float | None = None, reduction: NumericReductionT = 'mean', eps: float = 1e-12, return_weights: bool = False) → torch.Tensor | Tuple[torch.Tensor, torch.Tensor | None]

Pool neighbor-associated numeric values with optional weighting.

Parameters:

  • source_values – Numeric values associated with reference pool.

  • weighted – Whether to weight by similarity scores. Defaults to False.

  • scores – Optional custom scores for weighting. Defaults to None.

  • temperature – Optional temperature for softmax weighting. Defaults to None.

  • reduction – Reduction method (‘mean’, ‘weighted_mean’, ‘median’, ‘sum’, ‘max’, ‘min’). Defaults to ‘mean’.

  • eps – Numerical stability constant. Defaults to 1e-12.

  • return_weights – Whether to return applied weights. Defaults to False.

Returns:

Pooled values, and optionally the weights used.

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]

loki2.retrieve_from_sc.retrieve_with_celltype_filter(query_embeddings: torch.Tensor, embedding_pool: torch.Tensor, pool_labels: Sequence[Any] | numpy.ndarray | torch.Tensor, topk: int = 20, *, normalize_centroids: bool = True, return_assignments: bool = True) → RetrievalResult | Tuple[RetrievalResult, numpy.ndarray, torch.Tensor]

Assign queries to cell types via centroid similarity and retrieve neighbors only from matching cell types.

Parameters:

  • query_embeddings – (M, D) tensor of query embeddings.

  • embedding_pool – (N, D) tensor of reference embeddings.

  • pool_labels – Iterable of length N with cell type labels.

  • topk – Number of neighbors to retrieve per query.

  • normalize_centroids – Whether to L2-normalize cell type centroids.

  • return_assignments – If True, also return predicted labels and scores.

Returns:

RetrievalResult if return_assignments is False, otherwise a tuple of (RetrievalResult, predicted_labels, centroid_similarities).

loki2.retrieve_from_sc.usage() → None

Print usage information for the retrieval script.

loki2.retrieve_from_sc.sam
loki2.retrieve_from_sc.sample
loki2.retrieve_from_sc.epoch
loki2.retrieve_from_sc.output_dir
loki2.retrieve_from_sc.proj_path
loki2.retrieve_from_sc.sc_proj_path
loki2.retrieve_from_sc.sc_meta_path
loki2.retrieve_from_sc.output_dir
loki2.retrieve_from_sc.data_morph
loki2.retrieve_from_sc.sc_data_trans
loki2.retrieve_from_sc.adata_sc
loki2.retrieve_from_sc.query_embedding
loki2.retrieve_from_sc.embedding_pool
loki2.retrieve_from_sc.cell_ids_pool
loki2.retrieve_from_sc.pool_labels
loki2.retrieve_from_sc.retrieval_path
loki2.retrieve_from_sc.pred_scores_np
loki2.retrieve_from_sc.assign_df
loki2.retrieve_from_sc.assign_path