Einsum in python for a complex loop-CodePudding

I have complex loops in python that I'm trying to "vectorize" to improve computation time. I found the function np.einsum allowing it, I managed to use it, but I'm stuck with another loop.

In the following code, I put the loop I managed to "einsumize" (s1), and the other one where I didn't.

import numpy as np

Q = 6
P = 24
N = 40

bQ = np.arange(Q)
bP = np.arange(P)
uN = np.arange(N)
t1 = np.arange(P*Q*N).reshape([Q,P,N])
t2 = np.arange(Q*N*Q*N).reshape([Q,N,Q,N])

s1_ = np.einsum('p,q,n,qpn',bP, bQ, uN, t1)
s1 = 0
for p in range(P):
    for q in range(Q):
        for n in range(N):
            s1  = bP[p] * bQ[q] * uN[n] * t1[q,p,n]
print(s1)
print(s1_)
print()

s2_ = np.einsum('p,q,n,m,pnqm', bQ, bQ, uN, uN, t2)
s2 = 0
for p in range(Q):
    for q in range(Q):
        for n in range(N):
            for m in range(N):
                s2  = bQ[q] * bQ[q] * uN[n] * uN[m] * t2[p,n,q,m]
print(s2)
print(s2_)

The result of the previous code is

13475451600
13475451600

6125547636000
5707354770000

The math formula to compute s1 is : s1 = \sum_p\sum_q\sum_n bP[p] * bQ[q] * uN[n] * t1[q,p,n]. And the one to compute s2 is s2 = \sum_q\sum_q'\sum_n\sum_n' bQ[q] bQ[q'] uN[n] uN[n'] t2[q,n,q',n'].

For the triple loop, if I well understood how einsum works,I tell the different indices of the tensors that will be multiplied, and telling no output indices tell that all will be summed. But it seems not to be working for the quadruple loop.

EDIT : I saw an answer (which seems to have been deleted), telling that it was just a mistake of index on the quadruple loop... I should have seen it :/

CodePudding user response：

I see a typo in your code, you are not using the variable p outside of the order 4 tensor.

Try changing

                s2  = bQ[q] * bQ[q] * uN[n] * uN[m] * t2[p,n,q,m]

for

                s2  = bQ[p] * bQ[q] * uN[n] * uN[m] * t2[p,n,q,m]