沿 3D 数组轴的矢量化欧几里德距离 - Python

【问题标题】：Vectorized euclidean distance along an axis of a 3D array - Python沿 3D 数组轴的矢量化欧几里德距离 - Python
【发布时间】：2020-03-09 11:15:49
【问题描述】：

如何沿 3D 数组的轴应用对 2D 数组 (cdist) 进行操作的函数？

我尝试使用 numpy.apply_along_axis ，但我需要对 2D 数组而不是 1D 进行操作。我通过沿一个轴迭代获得了我需要的结果，但如果可能的话，我希望它矢量化：

from scipy import spatial
import numpy as np

a = np.random.randn(600).reshape(10, 20, 3)
distances = np.array([spatial.distance.cdist(a[i,:,:], a[i,:,:]) for i in range(a.shape[0])])

【问题讨论】：

你有一个 10x20 点的坐标系吗？你现在想计算这些点在 N=20 维度上的距离吗？（任何机会：D）
是的，你遇到了：D
很高兴 Divakar 的回答符合您的要求！

标签： python numpy scipy vectorization

【解决方案1】：

受this post的启发，我们可以用向量化的方式来解决它。因此，按照eucl_dist 包中的wiki contents（免责声明：我是它的作者），我们可以利用matrix-multiplication 和一些NumPy specific implementations，就像这样-

a_s = np.einsum('ijk,ijk->ij',a,a)
sq_dists = a_s[...,None]+a_s[...,None,:]-2*np.einsum('ijk,ilk->ijl',a,a)
dists = np.sqrt(sq_dists)

替代方案：

我们可以使用np.matmul/@-operator on Python3.x 来代替矩阵乘法部分。因此，np.einsum('ijk,ilk->ijl',a,a) 可以替换为 np.matmul(a,a.transpose(0,2,1))。

【讨论】：

做到了，非常感谢！我自己绝对猜不到。