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