Python中的欧几里得距离平方掩码向量化

Question

我正在运行代码，生成一个离A中某些位置距离小于D的位置的掩码。

N = [[0 for j in range(length_B)] for i in range(length_A)]    
dSquared = D*D

for i in range(length_A):
    for j in range(length_B):
        if ((A[j][0]-B[i][0])**2 + (A[j][1]-B[i][1])**2) <= dSquared:
            N[i][j] = 1

对于包含数万个位置的 A 和 B 列表，这段代码执行需要一些时间。我相信有一种向量化的方法可以使其运行更快。谢谢。

- ToneDaBass

---

这句话应该是“N [j] [i] = 1”而不是吗？ - Divakar

@Divakar - 你说得对，我为了在这里发布更简单的内容而更改了项目名称，结果自己也混淆了。我已经相应地更正了代码。 - ToneDaBass

Answer 1

使用二维数组索引更容易理解这段代码：

for j in range(length_A):
    for i in range(length_B):
        dist = (A[j,0] - B[i,0])**2 + (A[j,1] - B[i,1])**2 
        if dist <= dSquared:
            N[i, j] = 1

那个 dist 表达式看起来像：

((A[j,:] - B[i,:])**2).sum(axis=1)

如果只有2个元素，这个数组表达式可能不会更快，但它可以帮助我们重新思考问题。

我们可以使用广播来解决i,j和outter问题。

A[:,None,:] - B[None,:,:]  # 3d difference array

dist=((A[:,None,:] - B[None,:,:])**2).sum(axis=-1)  # (lengthA,lengthB) array

将其与 dSquared 进行比较，并使用生成的布尔数组作为在将 N 的元素设置为 1 时的掩码：

N = np.zeros((lengthA,lengthB))
N[dist <= dSquared] = 1

我没有测试过这段代码，所以可能存在漏洞，但我认为基本思路是正确的。这个想法可能足够让你解决其他情况的细节问题。

- hpaulj

Answer 2

您可以使用scipy的cdist进行距离计算，据说这种方法非常高效，例如：

from scipy.spatial.distance import cdist

N = (cdist(A,B,'sqeuclidean') <= dSquared).astype(int)

---

如@hpaulj的解决方案所建议的那样，我们也可以使用。现在，从问题中发布的代码来看，似乎我们正在处理Nx2形状的数组。因此，我们可以基本上切片第一列和第二列，并分别对它们执行广播减法。好处是我们不会进入3D，因此保持内存效率，这也可能转化为性能提升。因此，平方欧几里得距离将被计算如下 -

sq_eucl_dist = (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2

让我们计时这三种方法用于计算平方欧几里得距离的时间。

运行时间测试 -

In [75]: # Input arrays
    ...: A = np.random.rand(200,2)
    ...: B = np.random.rand(200,2)
    ...: 

In [76]: %timeit ((A[:,None,:] - B[None,:,:])**2).sum(axis=-1) # @hpaulj's solution
1000 loops, best of 3: 1.9 ms per loop

In [77]: %timeit (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2
1000 loops, best of 3: 401 µs per loop

In [78]: %timeit cdist(A,B,'sqeuclidean')
1000 loops, best of 3: 249 µs per loop

- Divakar

---

cdist.asType(bool) 与 np.nonzero() 结合使用，可以显著提高性能并减少内存占用，因为输出结果非常稀疏。感谢您的帮助，非常感激。 - ToneDaBass

Answer 3

我赞同以上使用Numpy的建议。循环代码也在对A进行比必要的索引操作。你可以尝试使用以下代码：

import numpy as np

dimension = 10000
A = np.random.rand(dimension, 2) + 0.0
B = np.random.rand(dimension, 2) + 1.0
N = []
d = 1.0

for i in range(len(A)):
    distances = np.linalg.norm(B - A[i,:], axis=1)
    for j in range(len(distances)):
        if distances[j] <= d:
            N.append((i,j))

print(len(N))

使用纯Python这样的方式很难获得良好的性能。我还要指出，多维数组的解决方案将需要...... 很多......内存。

- Steve Cohen

Answer 4

鉴于您的矩阵N可能是稀疏的，使用scipy.spatial.cKDTree会比基于暴力计算所有距离的任何方法具有更好的时间复杂度：

cKDTree(A).sparse_distance_matrix(cKDTree(B), max_distance=D)

- Eelco Hoogendoorn