Numpy：找到两个 3-D 数组之间的欧几里得距离

Posted 2023-02-23

【中文标题】Numpy：找到两个 3-D 数组之间的欧几里得距离

【英文标题】：Numpy: find the euclidean distance between two 3-D arrays

【发布时间】：2017-03-12 04:38:36

【问题描述】：

给定两个维度为 (2,2,2) 的 3-D 数组：

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

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

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

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

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

我尝试过 for 循环，但这些循环很慢，而且我正在按 (>>2, >>2, 2) 的顺序处理 3-D 数组。

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

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)]

【参考方案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 module 的另一种选择-

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

【讨论】：

您是否有链接或资源可以比较 np.einsum 与常规 numpy 效率？

@jadsq 如果我们谈论的是一般情况，我猜this one。在这个特定的案例中，我们使用einsum 一次性完成了两件事，这使得它在这里非常有用。

@Divakar，谢谢！今天我又从你那里学到了一些新东西。

numexpr 切片的速度要快得多。我的数组虽然很大：(8464, 8464, 2)，但我发现计算仍然太慢。

@user1658296 需要多少时间？

【参考方案2】：

为了完整性：

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

【讨论】：

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

【参考方案3】：

我建议在使用自定义平方和根而不是标准内置 math.hypot 和 np.hypot 时要格外小心。它们快速且经过优化且非常安全。

从这个意义上说np.linalg.norm(A-B, axis=-1) here 看起来最安全

对于非常大的矩阵，使用广播的 numpy 会受到内存限制并降低速度。在这种情况下，您可能希望使用 for 循环，但不要在速度上妥协。因为使用 numba 会很好here

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

参考：speed/overflow/underflow

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