计算numpy数组的某些条目之间的欧式距离

Question

I am a bit new to numpy and I am trying to calculate the pairwaise distance between some of the elements of a numpy array.

I have a numpy nx 3 array with n 3D cartesian coordinates (x,y,z) representing particles in a grid. Some of these particles move as the program runs and I need to keep track of the distances of the ones that move. I hold a list of integers with the index of the particles that have moved.

I am aware of pdist but this calculates the distance between every pair of particles, which would be inefficient as only some of them have moved. Ideally, for example, if only 1,2 have moved then I would only calculate the distance of 1 with 2...N and 2 with 3...N

What would be the most efficient way of doing this? Right now I have a double loop which doesn't seem ideal...

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

Thanks

Answer 1

I believe this is what you want (create new axis and use broadcasting for full vectorization):

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

If you wan't to use a jit compiler, here is an example using 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

Performance compared to Julien's solution

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

This gives on my Core i7-4771 0.15s for the Numba solution and 2.4s for Julien's solution.