计算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.15 秒的 Numba 解决方案和 2.4 秒的 Julien 解决方案。