如何在numpy数组上进行nD距离和最近邻居计算

Question

此问题旨在成为规范的重复目标

给定两个数组X和Y形状,(i, n)并(j, n)表示 - n维坐标列表,

def test_data(n, i, j, r = 100):
    X = np.random.rand(i, n) * r - r / 2
    Y = np.random.rand(j, n) * r - r / 2
    return X, Y

X, Y = test_data(3, 1000, 1000)

找到最快的方法是什么:

1. 每个点和每个点之间的D形状距离(i,j)XY
2. 该指数k_i与距离k_d的的k针对所有点最近的邻居X中的每一个点Y
3. 该指数r_i,r_j和距离r_d的每一个点在X距离之内r的每一个点的j中Y

鉴于以下几组限制:

• 只使用 numpy
• 使用任何python包

包括特例:

• Y 是 X

在所有情况下,距离主要表示欧几里德距离,但可以随意突出显示允许其他距离计算的方法.

Answer 1

1.所有距离

• 只使用 numpy

天真的方法是:

D = np.sqrt(np.sum((X[:, None, :] - Y[None, :, :])**2, axis = -1))

然而,这会占用大量内存来创建一个(i, j, n)形状中间矩阵,并且非常慢

但是,由于来自@Divakar(eucl_dist包,wiki)的技巧,我们可以使用一些代数并np.einsum进行分解: (X - Y)**2 = X**2 - 2*X*Y + Y**2

D = np.sqrt(                                #  (X - Y) ** 2   
np.einsum('ij, ij ->i', X, X)[:, None] +    # = X ** 2        \
np.einsum('ij, ij ->i', Y, Y)          -    # + Y ** 2        \
2 * X.dot(Y.T))                             # - 2 * X * Y

• Y 是 X

与上面类似:

XX = np.einsum('ij, ij ->i', X, X)
D = np.sqrt(XX[:, None] + XX - 2 * X.dot(X.T))

请注意,使用此方法,浮点不精确可以使对角线项与零略微偏离.如果您需要确保它们为零,则需要明确设置它:

np.einsum('ii->i', D)[:] = 0 

• 任何包裹

scipy.spatial.distance.cdist 是最直观的内置功能,比裸机快得多 numpy

from scipy.spatial.distance import cdist
D = cdist(X, Y)

cdist还可以处理许多距离测量以及用户定义的距离测量(尽管这些测量未被优化).查看上面链接的文档了解详细信息.

• Y 是 X

对于自引用距离,scipy.spatial.distance.pdist工作类似于cdist,但返回1-D压缩距离数组,通过仅使每个项一次来节省对称距离矩阵上的空间.您可以使用将其转换为方阵squareform

from scipy.spatial.distance import pdist, squareform
D_cond = pdist(X)
D = squareform(D_cond)

2. K最近邻居(KNN)

• 只使用 numpy

我们可以np.argpartition用来获取k-nearest索引并使用它们来获得相应的距离值.因此,D当数组保持上面获得的距离值时,我们将 -

if k == 1:
    k_i = D.argmin(0)
else:
    k_i = D.argpartition(k, axis = 0)[:k]
k_d = np.take_along_axis(D, k_i, axis = 0)

然而,在我们减少数据集之前,我们可以通过不取平方根来加快这一点. np.sqrt是计算欧几里德范数最慢的部分,所以我们不希望这样做直到结束.

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +\
       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)
if k == 1:
    k_i = D_sq.argmin(0)
else:
    k_i = D_sq.argpartition(k, axis = 0)[:k]
k_d = np.sqrt(np.take_along_axis(D_sq, k_i, axis = 0))

现在,np.argpartition执行间接分区,并不一定按排序顺序给我们元素,只确保第一个k元素是最小元素.因此,对于排序输出,我们需要使用argsort上一步的输出 -

sorted_idx = k_d.argsort(axis = 0)
k_i_sorted = np.take_along_axis(k_i, sorted_idx, axis = 0)
k_d_sorted = np.take_along_axis(k_d, sorted_idx, axis = 0)

• 任何包裹

KD-Tree是一种快速查找邻居和约束距离的方法.最推荐的方法是使用k_i's X或Y

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +\
       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)
if k == 1:
    k_i = D_sq.argmin(0)
else:
    k_i = D_sq.argpartition(k, axis = 0)[:k]
k_d_sq = np.take_along_axis(D_sq, k_i, axis = 0)
sorted_idx = k_d_sq.argsort(axis = 0)
k_i_sorted = np.take_along_axis(k_i, sorted_idx, axis = 0)

不幸n的是,KDTree的实现速度很慢,并且存在对较大数据集进行分段错误的倾向.正如@HansMusgrave 在这里指出的那样,2**n增加了很多性能,但并不像普通的那样常见,scipy并且目前只能处理欧几里德距离(而scipy.spatial.KDTreein scipy.spatial.cKDTree可以处理任何顺序的Minkowsi p规范)

• 任意指标

BallTree具有与KDTree类似的算法属性.我不知道Python中的并行/矢量化/快速BallTree,但是使用scipy我们仍然可以对用户定义的指标进行合理的KNN查询.如果可用,内置指标将更快.

D_sq = np.einsum('ij, ij ->i', X, X)[:, None] +\
       np.einsum('ij, ij ->i', Y, Y) - 2 * X.dot(Y.T)

如果不是指标,这个答案将是错误的.BallTree比蛮力更快的唯一原因是因为指标的属性允许它排除某些解决方案.对于真正的任意功能,蛮力实际上是必要的.scipy

3.半径搜索

• 只使用 pykdtree

最简单的方法就是使用布尔索引:

XX = np.einsum('ij, ij ->i', X, X)
D_sq = XX[:, None] + XX - 2 * X.dot(X.T))

• 任何包裹

与上面类似,你可以使用 scipy

from scipy.spatial.distance import KDTree
X_tree = KDTree(X)
k_d, k_i = X_tree.query(Y, k = k)

要么 KDTree

k_d, k_i = X_tree.query(X, k = k)

不幸的是,scipy最终成为一个索引数组列表,有点难以解决以供以后使用.

更容易使用X's Y,它可以输出一个d()

def d(a, b):
    return max(np.abs(a-b))

tree = sklearn.neighbors.BallTree(X, metric=d)
k_d, k_i = tree.query(Y)

这是距离矩阵的一种非常灵活的格式,因为它保留了一个实际的矩阵(如果转换为numpy)也可以用于许多矢量化操作.