Question

这段代码片段是我的一个项目的瓶颈。是否有函数调用可以替换for循环并加速它？

D = np.zeros((nOcc,nOcc,nVir,nVir))
for i in range(nOcc):
   for j in range(i+1):
      tmp = Ew[i] + Ew[j]
      for a in range(nVir):
         tmp2 = tmp - Ew[a+nOcc]
         for b in range(a+1):
            tmp3 = 1.0/(tmp2 - Ew[b+nOcc])
            D[i,j,a,b] = Iiajb[i,a,j,b]*tmp3
            D[i,j,b,a] = Iiajb[i,b,j,a]*tmp3
            D[j,i,a,b] = D[i,j,b,a]
            D[j,i,b,a] = D[i,j,a,b]

Answer 1

首先，让我们生成一些任意数据，这符合一些必要的原则：

nOcc = 30
nVir = 120
Ew = np.random.rand(nOcc+nVir)
Ew[:nOcc]*=-1
Ia = np.random.rand(nOcc)
Ib = np.random.rand(nVir)
I = np.einsum('a,b,c,d->abcd',Ia,Ib,Ia,Ib)

让我们将您的基本代码包装为例：

def oldcalc_D(Iiajb,nOcc,nVir,Ew):
    D = np.zeros((nOcc,nOcc,nVir,nVir))
    for i in range(nOcc):
       for j in range(i+1):
          tmp = Ew[i] + Ew[j]
          for a in range(nVir):
             tmp2 = tmp - Ew[a+nOcc]
             for b in range(a+1):
                tmp3 = 1.0/(tmp2 - Ew[b+nOcc])
                D[i,j,a,b] = Iiajb[i,a,j,b]*tmp3
                D[i,j,b,a] = Iiajb[i,b,j,a]*tmp3
                D[j,i,a,b] = D[i,j,b,a]
                D[j,i,b,a] = D[i,j,a,b]
    return D

利用整体对称性通常是一种很好的策略;然而，单独使用numpy并不值得花费，所以我们将忽略这一点并简单地对代码进行矢量化：

def newcalc_D(I,nOcc,nVir,Ew):
    O = Ew[:nOcc]
    V = Ew[nOcc:]
    D = O[:,None,None,None] - V[:,None,None] + O[:,None] - V
    return (I/D).swapaxes(1,2)

一些时间：

np.allclose(oldcalc_D(I,nOcc,nVir,Ew),newcalc_D(I,nOcc,nVir,Ew))
True

%timeit newcalc_D(I,nOcc,nVir,Ew)
1 loops, best of 3: 142 ms per loop

%timeit oldcalc_D(I,nOcc,nVir,Ew)
1 loops, best of 3: 15 s per loop

所以只有约100倍的速度，因为我说这是一个相当简单的传递，让你知道该怎么做。这可以做得更好，但应该是计算的一个微不足道的部分，因为积分变换是（O）N ^ 5对比这在（O）N ^ 4。对于这些操作，我使用numba的autojit功能：

from numba import autojit

numba_D = autojit(oldcalc_D)

%timeit numba_D(I,nOcc,nVir,Ew)
10 loops, best of 3: 55.1 ms per loop

Answer 2

除非这是Python3，否则您可能需要先将range替换为xrange：前者创建整个列表，而后者只是一个迭代器，在这种情况下您只需要它。
对于较大的N s，速度差异应该是显而易见的。

此外，看作你的使用numpy，可能有一种矢量化的方式来实现该算法。如果是这种情况，矢量化实现应该快几个数量级。但除非你解释变量和算法，否则我们无法帮助你朝这个方向发展。

是否有函数调用可以替换此代码中的for循环？

2 个答案: