沿着3D阵列的轴的向量化欧式距离-Python

Question

如何在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])])

Answer 1

受的启发this post，我们可以采用向量化的方式解决该问题。因此，继 wiki contents从eucl_dist 包（免责声明：我是它的作者），我们可以利用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))。