计算numpy数组的一些条目之间的欧氏距离

Question

我对numpy有点新意，我正在尝试计算numpy数组的某些元素之间的成对距离。

我有一个numpy n x 3数组，其中n个3D笛卡尔坐标（x，y，z）代表网格中的粒子。这些粒子中的一些随着程序的运行而移动，我需要跟踪移动的粒子的距离。我持有一个整数列表，其中包含已移动的粒子的索引。

我知道pdist，但是这会计算每对粒子之间的距离，因为只有一些粒子移动了，所以效率很低。理想情况下，例如，如果只有1,2移动了那么我只计算1的距离为2 ... N和2与3 ... N

这样做最有效的方法是什么？现在我有一个看起来不太理想的双循环......

for i in np.nditer(particles_moved):
    particles = particles[particles!=i]
    for j in np.nditer(particles):
        distance(xyz,i, j)

谢谢

Answer 1

我相信这就是你想要的（创建新轴并使用广播进行完全矢量化）：

import numpy as np

particles = np.arange(12).reshape((-1,3))
moved = np.array([0,2])
np.linalg.norm(particles[moved][:,None,:]-particles[None,:,:], axis=-1)

array([[  0.        ,   5.19615242,  10.39230485,  15.58845727],
       [ 10.39230485,   5.19615242,   0.        ,   5.19615242]])

Answer 2

如果你不想使用jit编译器，这里有一个使用Numba的例子

@nb.njit(fastmath=True,parallel=True)
def distance(Paticle_coords,indices_moved):
  dist_res=np.empty((indices_moved.shape[0],Paticle_coords.shape[0]),dtype=Paticle_coords.dtype)

  for i in range(indices_moved.shape[0]):
    Koord=Paticle_coords[indices_moved[i],:]
    dist_res[i,:]=np.sqrt((Koord[0]-Paticle_coords[:,0])**2+(Koord[1]-Paticle_coords[:,1])**2+(Koord[2]-Paticle_coords[:,2])**2)

  return dist_res

性能与Julien的解决方案相比

#Create Data
Paticle_coords=np.random.rand(10000000,3)
indices_moved=np.array([0,5,6,3,7],dtype=np.int64)

这使我的Core i7-4771 0.15s用于Numba解决方案，而2.4s用于Julien的解决方案。