# Optimizing for loops by using Einstein summations instead

I have some Python code containing a couple of for loops, which I would like to optimize by using low-level functions. My first approach is to use numpy.einsum(), but the results don't match the for-loop implementation:

## first term
T11 = 0.125 * numpy.einsum('abcd,efab,cdef,abcd,efcd',
g[:Nelec,:Nelec,Nelec:,Nelec:], # <ab||rs> -> abcd
g[:Nelec,:Nelec,:Nelec,:Nelec], # <cd||ab> -> efab
g[Nelec:,Nelec:,:Nelec,:Nelec], # <rs||cd> -> cdef
D,                              # e_a + e_b - e_r - e_s -> abcd
D,                              # e_c + e_d - e_r - e_s -> efcd
)

T12 = 0
for a in range(Nelec):
for b in range(Nelec):

for c in range(Nelec):
for d in range(Nelec):

for r in range(Nvirt):
for s in range(Nvirt):

# nominator
nom  = g[a,b,Nelec+r,Nelec+s]
nom *= g[c,d,a,b]
nom *= g[Nelec+r,Nelec+s,c,d]

# denominator
denom = D[a,b,r,s] * D[c,d,r,s]

T12 += nom*denom

T12 *= 0.125

## second term
T21 = 0.125 * numpy.einsum('abcd,cdef,efab,abcd,abef',
g[:Nelec,:Nelec,Nelec:,Nelec:], # <ab||rs> -> abcd
g[Nelec:,Nelec:,Nelec:,Nelec:], # <rs||tu> -> cdef
g[Nelec:,Nelec:,:Nelec,:Nelec], # <tu||ab> -> efab
D,                              # e_a + e_b - e_r - e_s -> abcd
D,                              # e_a + e_b - e_t - e_u -> abef
)

T22 = 0
for a in range(Nelec):
for b in range(Nelec):

for r in range(Nvirt):
for s in range(Nvirt):

for t in range(Nvirt):
for u in range(Nvirt):

# nominator
nom  = g[a,b,Nelec+r,Nelec+s]
nom *= g[Nelec+r,Nelec+s,Nelec+t,Nelec+u]
nom *= g[Nelec+t,Nelec+u,a,b]

# denominator
denom = D[a,b,r,s] * D[a,b,t,u]

T22 += nom*denom

T22 *= 0.125

## third term
T31 = numpy.einsum('abcd,edfb,cfae,abcd,aecf',
g[:Nelec,:Nelec,Nelec:,Nelec:], # <ab||rs> -> abcd
g[:Nelec,Nelec:,Nelec:,:Nelec], # <cs||tb> -> edfb
g[Nelec:,Nelec:,:Nelec,:Nelec], # <rt||ac> -> cfae
D,                              # e_a + e_b - e_r - e_s -> abcd
D,                              # e_a + e_c - e_r - e_t -> aecf
)

T32 = 0
for a in range(Nelec):
for b in range(Nelec):
for c in range(Nelec):
for r in range(Nvirt):
for s in range(Nvirt):
for t in range(Nvirt):

# nominator
nom  = g[a,b,Nelec+r,Nelec+s]
nom *= g[c,Nelec+s,Nelec+t,b]
nom *= g[Nelec+r,Nelec+t,a,c]

# denominator
denom = D[a,b,r,s] * D[a,c,r,t]

T32 += nom*denom

print(T11, T12)
print(T21, T22)
print(T31, T32)
print(T11+T21+T31, T12+T22+T32)


For my test case I get the following output:

 0.016038311274  0.016038311274
0.014256694984  0.014256694984
-0.044483006105 -0.044483006105
-0.014187999845 -0.014187999845


The first column contains the einsum() results, the second column is for the correctly working for loop implementation.

So my question is: Why does einsum() not give the same results (except for the second term $T_2$)? Is there some error in my einsum() implementation or is there a conceptual problem here?

Edit: After fixing one error in T11 and another error in T31 both methods yield same results.

Not necessarily important, but the equations are for the Møller-Plesset perturbation theory third-order energy correction:

$$E_0^{(3)} = T_1 + T_2 + T_3 \\ T_1 = \frac{1}{8} \sum_{abcdrs} \frac{\langle ab||rs\rangle\langle cd||ab\rangle\langle rs||cd\rangle}{(\epsilon_a+\epsilon_b-\epsilon_r-\epsilon_s)(\epsilon_c+\epsilon_d-\epsilon_r-\epsilon_s)} \\ T_2 = \frac{1}{8} \sum_{abrstu} \frac{\langle ab||rs\rangle\langle rs||tu\rangle\langle tu||ab\rangle}{(\epsilon_a+\epsilon_b-\epsilon_r-\epsilon_s)(\epsilon_c+\epsilon_d-\epsilon_t-\epsilon_u)} \\ T_3 = \sum_{abcrst} \frac{\langle ab||rs\rangle\langle cs||tb\rangle\langle rt||ac\rangle}{(\epsilon_a+\epsilon_b-\epsilon_r-\epsilon_s)(\epsilon_a+\epsilon_c-\epsilon_r-\epsilon_t)}$$

In the code I use:

• g[i,j,k,l] = $\langle ij||kl\rangle$ being the anti-symmetrized two electron repulsion integrals
• D[a,b,r,s] = $\frac{1}{\epsilon_i+\epsilon_j-\epsilon_r-\epsilon_s}$ with $\epsilon_i$ being the energy of orbital $i$
• Nelec for the number of electrons, which equals the number of occupied spin orbitals
• Norb for the total number of spin orbitals
• Nvirt = Norb - Nelec for the number of virtual (unoccupied) orbitals
• orbitals are ordered by energy, thus occupied first, then virtuals

And orbital indices follow the convention:

• a,b,c,d running over occupied orbitals only
• r,s,t,u running over virtual orbitals only
• i,j,k,l running over all orbitals
• I'm not familiar with the operation of numpy's einsum, but generically one wouldn't expect the same index to appear three times in Einstein notation. – origimbo Jul 12 '17 at 22:55
• I tried to find something about this in the numpy documentation. But it only talks about "repeated" indices, there is no statement on the amount of repetitions. – Feodoran Jul 13 '17 at 6:43
• The summation is expressed incorrectly for the numerator of $T_{3}$. You repeated the same one as for $T_{2}$. I'm guessing the string for einsum should be 'abcd,edfb,cfae,abcd,abef' – gpavanb Jul 20 '17 at 2:07
• Thank you, this was probably a copy-paste-error. I found another error in the code for T11, after fixing both, the results match. Therefore, einsum works for indices repeating more than twice as well. However, a straight-forward Cython implementation is far more efficient (factor of 10 for my test case). – Feodoran Jul 20 '17 at 8:59