コード例 #1
0
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
コード例 #2
0
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