在 Python 中使用矩阵向量乘法进行速度优化

Question

我想在python中优化以下代码：

for imode in N.arange(3*natom): #Loop on perturbation (6 for 2 atoms)
  for ikpt in N.arange(nkpt):
    for iband in N.arange(nband):    
      for iatom1 in N.arange(natom):
        for iatom2 in N.arange(natom):
          for idir1 in N.arange(0,3):
            for idir2 in N.arange(0,3):    
              fan_corrQ[imode,ikpt,iband] += EIG2D[ikpt,iband,idir1,iatom1,idir2,iatom2]*\
                  displ_red_FAN2[imode,iatom1,iatom2,idir1,idir2]
              ddw_corrQ[imode,ikpt,iband] += ddw_save[ikpt,iband,idir1,iatom1,idir2,iatom2]*\
                  displ_red_DDW2[imode,iatom1,iatom2,idir1,idir2]

如您所见，我想对我的多空间 python 数组的一些索引进行求和。 我想要类似的东西：

for imode in N.arange(3*natom): #Loop on perturbation (6 for 2 atoms)
  for ikpt in N.arange(nkpt):
    for iband in N.arange(nband):
      fan_corrQ[imode,ikpt,iband] = N.dot(EIG2D[ikpt,iband,:,:,:,:],displ_red_FAN2.T[imode,:,:,:,:])
      ddw_corrQ[imode,ikpt,iband] = N.dot(ddw_save[ikpt,iband,:,:,:,:],displ_red_DDW2.T[imode,:,:,:,:])

当然，我有一个问题是不能乘以相同的索引，所以我重新定义了它。我还必须指出，我正在处理复数：

displ_red_DDW2 = N.zeros((3*natom,3,natom,3,natom),dtype=complex)

我还在 ipython 中尝试了一个小的虚拟程序来测试它：

import numpy as N
atom =2
displ_red_FAN2 = N.zeros((3*natom,3,natom,3,natom),dtype=complex)
EIG2D = N.zeros((216,12,3,2,3,2))

所以我有 displ_red_FAN2.shape = (6, 3, 2, 3, 2) 和 EIG2D.shape = (216, 12, 3, 2, 3, 2)

所以如果我这样做：

N.dot(EIG2D[1,1,:,:,:,:],displ_red_FAN2[1,:,:,:,:].T).shape

它应该给出 (3,2,3,2) 但它给出的是 (3, 2, 3, 2, 3, 3) ???然后当乘法完成时，我想我将不得不做一些求和来降低维度。

任何帮助都会很棒！

干杯！

塞缪尔

Answer 1

你可以使用np.tensordot()

fan_corrQ = np.tensordot(displ_red_FAN2, EIG2D, axes = ([3,1,4,2],[2,3,4,5]))
ddw_corrQ = np.tensordot(displ_red_DDW2, ddw_save, axes = ([3,1,4,2],[2,3,4,5]))

这给出了与您当前方法相同的结果，并且速度大约为 9 times。

关于您的其他问题。 np.dot() 对应于 ND-array 的作用是 summing over the last axis of the first array and the second-to-last axis of the second array。

Answer 2

我认为您可以通过使用 Einstein 求和 (np.einsum) 大大简化这一过程。语法可能有点难以理解，所以我稍微简化了变量和索引的名称：

# arrays
EIG2D           --> A
displ_red_FAN2  --> B
fan_corrQ       --> C

# indices
ikpt    --> i
iband   --> j
idir1   --> k
iatom1  --> l
idir2   --> m
iatom2  --> n
imode   --> o

np.einsum 采用逗号分隔的下标列表，每个下标指代 对应输入数组的维度。每当索引重复时，它 在输出中求和。您还可以在 通过给出输出索引来输出。

在你的情况下，我认为：

...
fan_corrQ[imode,ikpt,iband] += EIG2D[ikpt,iband,idir1,iatom1,idir2,iatom2]*\
    displ_red_FAN2[imode,iatom1,iatom2,idir1,idir2]
...

应该简化为：

C = np.einsum('ijklmn,olnkm->oij',A,B)

您应该尝试一下，并确保我没有犯任何错误！与sgpc 的回答一样，在内存要求方面也有类似的注意事项。