結合を使用してnumpyでテンソルをロードする並列化

Question

特定の問題を解決するために物理学で次のコードが使用されています。それは動作しますが、非常に遅く、私はそれが最適化できると信じています。 （しかし、根本的な違いを伴う）非常によく似たケースがここに示されている：parallelize (not symmetric) loops in python

G_tensor = numpy.matlib.identity(N_particles*3,dtype=complex) 

for i in range(N_particles): 
    for j in range(N_particles): 
     if i != j: 

      #Do lots of things, here is shown an example. 
      # However you should not be scared because 
      #it only fills the G_tensor 
      R = numpy.linalg.norm(numpy.array(positions[i])-numpy.array(positions[j])) 
      rx = numpy.array(positions[i][0])-numpy.array(positions[j][0]) 
      ry = numpy.array(positions[i][1])-numpy.array(positions[j][1]) 
      rz = numpy.array(positions[i][2])-numpy.array(positions[j][2]) 

      pf = -numpy.exp(1j*k*R)/(4*math.pi*R) 
      b = (k/R)*(1j*k*R-1.)/(k*R) 
      G_tensor[3*i+0,3*j+0] = 0 #Gxx 
      G_tensor[3*i+1,3*j+1] = 0 #Gyy 
      G_tensor[3*i+2,3*j+2] = 0 #Gzz 
      G_tensor[3*i+0,3*j+1] = pf*(b * (-rz)/R)  #Gxy 
      G_tensor[3*i+0,3*j+2] = pf*(b * (ry)/R)  #Gxz 
      G_tensor[3*i+1,3*j+0] = pf*(b * (rz)/R)  #Gyx 
      G_tensor[3*i+1,3*j+2] = pf*(b * (-rx)/R)  #Gyz 
      G_tensor[3*i+2,3*j+0] = pf*(b * (-ry)/R)  #Gzx 
      G_tensor[3*i+2,3*j+1] = pf*(b * (rx)/R)  #Gzy 

それは彼がコードを最適化するためにnumpyのintrisic機能を使用@jadsqによって与えられるいずれかのような解決策を与えることができます。

Answer 1

# this particular shape correctly reshapes to (3N,3N) without copy 
G_tensor = np.zeros((N, 3, N, 3), dtype=complex) 

r = positions[:,None] - positions 
R = np.linalg.norm(r, axis=-1) 
np.fill_diagonal(R, 1.0) # prevent divide by zero 

pf = -np.exp(1j*k*R)/(4*np.pi*R) 
b = (k/R)*(1j*k*R-1.)/(k*R) 

r = r.transpose(2, 0, 1) 
G_tensor[:,[2,0,1],:,[1,2,0]] = pf*(b * (r)/R) 
G_tensor[:,[1,2,0],:,[2,0,1]] = pf*(b * (-r)/R) 
G_tensor[np.arange(N),:,np.arange(N),:] = np.eye(3) 
G_tensor.shape = (3*N, 3*N) 

"Combining advanced and basic indexing"のドキュメントも参照してください。