def create_tree(data, params):
'''
Create a faiss/cKDTree/KDTree/annoy/pynndescent index for nearest neighbour lookup.
All undescribed input as in ``bbknn.bbknn()``. Returns the resulting index.
Input
-----
data : ``numpy.array``
PCA coordinates of a batch's cells to index.
params : ``dict``
A dictionary of arguments used to call ``bbknn.matrix.bbknn()``, plus ['computation']
storing the knn algorithm to use.
'''
if params['computation'] == 'annoy':
ckd = AnnoyIndex(data.shape[1], metric=params['metric'])
for i in np.arange(data.shape[0]):
ckd.add_item(i, data[i, :])
ckd.build(params['annoy_n_trees'])
elif params['computation'] == 'pynndescent':
ckd = pynndescent.NNDescent(
data,
metric=params['metric'],
n_jobs=-1,
n_neighbors=params['pynndescent_n_neighbors'],
random_state=params['pynndescent_random_state'])
ckd.prepare()
elif params['computation'] == 'faiss':
ckd = faiss.IndexFlatL2(data.shape[1])
ckd.add(data)
elif params['computation'] == 'cKDTree':
ckd = cKDTree(data)
elif params['computation'] == 'KDTree':
ckd = KDTree(data, metric=params['metric'])
return ckd
def create_tree(data,approx,metric,use_faiss,n_trees): ''' Create a faiss/cKDTree/KDTree/annoy index for nearest neighbour lookup. All undescribed input as in ``bbknn.bbknn()``. Returns the resulting index. Input ----- data : ``numppy.array`` PCA coordinates of a batch's cells to index. ''' if approx: ckd = AnnoyIndex(data.shape[1],metric=metric) for i in np.arange(data.shape[0]): ckd.add_item(i,data[i,:]) ckd.build(n_trees) elif metric == 'euclidean': if 'faiss' in sys.modules and use_faiss: ckd = faiss.IndexFlatL2(data.shape[1]) ckd.add(data) else: ckd = cKDTree(data) else: ckd = KDTree(data,metric=metric) return ckd
