Numpy：找出两个三维数组之间的欧几里得距离

Question

假设有两个维度为(2,2,2)的三维数组:

A = [[[ 0,  0],
    [92, 92]],

   [[ 0, 92],
    [ 0, 92]]]

B = [[[ 0,  0],
    [92,  0]],

   [[ 0, 92],
    [92, 92]]]

如何高效地为A和B中的每个向量找到欧几里得距离？

我尝试了for循环，但这些循环速度较慢，并且我正在处理大小为（>>2，>>2，2）的三维数组。

最终，我想要一个如下形式的矩阵：

C = [[d1, d2],
     [d3, d4]]

编辑：

我尝试了以下循环，但最大的问题是它会丢失我想要保留的维度。但距离是正确的。

[numpy.sqrt((A[row, col][0] - B[row, col][0])**2 + (B[row, col][1] -A[row, col][1])**2) for row in range(2) for col in range(2)]

- user1658296

Answer 1

以NumPy向量化的方式思考，执行逐元素微分、平方和最后沿着最后一个轴求平方根。因此，直接的实现方案是 -

np.sqrt(((A - B)**2).sum(-1))

我们可以使用np.einsum一次性完成沿着最后一个轴的平方和运算，从而使其更加高效，代码如下 -

subs = A - B
out = np.sqrt(np.einsum('ijk,ijk->ij',subs,subs))

使用numexpr模块的另一种选择 -

import numexpr as ne
np.sqrt(ne.evaluate('sum((A-B)**2,2)'))

由于我们在最后一个轴上使用了长度为2，因此我们可以只对其进行切片并将其馈送到evaluate方法中。请注意，在评估字符串内部无法进行切片。因此，修改后的实现如下：

a0 = A[...,0]
a1 = A[...,1]
b0 = B[...,0]
b1 = B[...,1]
out = ne.evaluate('sqrt((a0-b0)**2 + (a1-b1)**2)')

运行时测试

函数定义 -

def sqrt_sum_sq_based(A,B):
    return np.sqrt(((A - B)**2).sum(-1))

def einsum_based(A,B):
    subs = A - B
    return np.sqrt(np.einsum('ijk,ijk->ij',subs,subs))

def numexpr_based(A,B):
    return np.sqrt(ne.evaluate('sum((A-B)**2,2)'))

def numexpr_based_with_slicing(A,B):
    a0 = A[...,0]
    a1 = A[...,1]
    b0 = B[...,0]
    b1 = B[...,1]
    return ne.evaluate('sqrt((a0-b0)**2 + (a1-b1)**2)')

时间 -

In [288]: # Setup input arrays
     ...: dim = 2
     ...: N = 1000
     ...: A = np.random.rand(N,N,dim)
     ...: B = np.random.rand(N,N,dim)
     ...: 

In [289]: %timeit sqrt_sum_sq_based(A,B)
10 loops, best of 3: 40.9 ms per loop

In [290]: %timeit einsum_based(A,B)
10 loops, best of 3: 22.9 ms per loop

In [291]: %timeit numexpr_based(A,B)
10 loops, best of 3: 18.7 ms per loop

In [292]: %timeit numexpr_based_with_slicing(A,B)
100 loops, best of 3: 8.23 ms per loop

In [293]: %timeit np.linalg.norm(A-B, axis=-1) #@dnalow's soln
10 loops, best of 3: 45 ms per loop

- Divakar

---

你是否有比较np.einsum和普通numpy效率的链接或资源？ - jadsq

@jadsq 我猜这个如果我们谈论一般情况。在这个特定的情况下，我们使用einsum一次完成两件事情，这使得它在这里非常有益。 - Divakar

@Divakar，谢谢！我今天又从你这里学到了新的东西。 - MaxU - stand with Ukraine

numexpr在切片操作时速度更快。但是我的数组非常大：（8464，8464，2），我发现计算仍然太慢了。 - user1658296

我得到了大约2.5秒的时间，比linalg要好得多，后者平均需要大约13秒。 - user1658296

Answer 2

为了完整性：

np.linalg.norm(A-B, axis=-1)

- dnalow

---

不错！已将此添加到运行时测试方法列表中。 - Divakar

Answer 3

当使用自定义平方和根而非标准内置的math.hypot和np.hypot时，我建议非常小心。后者速度快、经过优化且很安全。

从这个角度来看，np.linalg.norm(A-B, axis=-1)在这里看起来最安全

对于非常大的矩阵，numpy使用广播将会受到内存限制并变慢。在这种情况下，你需要使用for循环，但不要妥协速度。为此使用numba是好的在这里

i, j, k = 1e+200, 1e+200, 1e+200
math.hypot(i, j, k)
# 1.7320508075688773e+200

参考：速度/溢出/下溢

np.sqrt(np.sum((np.array([i, j, k])) ** 2))
# RuntimeWarning: overflow encountered in square

- eroot163pi