scimilarity.cell_search_knn#

class scimilarity.cell_search_knn.CellSearchKNN(model_path, knn_type, use_gpu=False)[source]#

Bases: CellEmbedding

A class for searching similar cells using cell embeddings kNN.

Parameters:

  • model_path (str) – Path to the directory containing model files.

  • knn_type (str, default: "hnswlib") – What type of knn to use, options are [“hnswlib”, “tiledb_vector_search”]

  • use_gpu (bool, default: False) – Use GPU instead of CPU.

Examples

>>> cs = CellSearchKNN(model_path="/opt/data/model")
get_nearest_neighbors(embeddings, k=50, ef=100)[source]#

Get nearest neighbors. Used by classes that inherit from CellEmbedding and have an instantiated kNN.

Parameters:

  • embeddings (numpy.ndarray) – Embeddings as a 2D numpy array.

  • k (int, default: 50) – The number of nearest neighbors.

  • ef (int, default: 100) – The size of the dynamic list for the nearest neighbors for hnswlib. See nmslib/hnswlib

Returns:

  • nn_idxs (numpy.ndarray) – A 2D numpy array of nearest neighbor indices [num_embeddings x k].

  • nn_dists (numpy.ndarray) – A 2D numpy array of nearest neighbor distances [num_embeddings x k].

Return type:

Tuple[numpy.ndarray, numpy.ndarray]

Examples

>>> nn_idxs, nn_dists = get_nearest_neighbors(embeddings)
load_knn_index(knn_file, memory_budget=50000000)[source]#

Load the kNN index file

Parameters:

  • knn_file (str) – Filename of the kNN index.

  • memory_budget (int, default: 50000000) – Memory budget for tiledb vector search.