python - 在 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.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 的回答，类似的警告适用于内存要求。