Making a np.einsum faster when inputs are many identical arrays? (Or any other faster method)

user_9215333 0
Viewed 79

I have a piece of code of type:

nnt = np.real(np.einsum('xa,xb,yc,yd,abcde->exy',evec,evec,evec,evec,quartic))

where evec is (say) an L x L np.float32 array, and quartic is a L x L x L x L x T np.complex64 array.

I found that this routine is rather slow.

I thought that since all the evec's are identical, there might be a faster way of doing this?

Thanks in advance.

numpy-einsum numpy python
1 Answers

For start you can reuse the first calculation:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic))

And if you don't care about memory and only need performance:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic,optimize=True))
user_49759810
Related