使用组合加载张量与 numpy 的并行化答案

【问题标题】：Parallelization of loading tensor with numpy using combinatory使用组合加载张量与 numpy 的并行化
【发布时间】：2017-01-19 14:55:58
【问题描述】：

以下代码在物理学中用于解决特定问题。它可以工作，但是速度很慢，我相信可以优化。此处显示了一个非常相似的案例（但有根本的区别）：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

是否有可能给出像@jadsq 给出的那样的解决方案，他使用numpy intrisic 函数来优化代码。

【问题讨论】：

标签： python numpy optimization parallel-processing

【解决方案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" 的文档。

【讨论】：