CODEWITHSUNDEEP
pythonnumpynumpy-einsum
purestate
purestate

Reputation: 63

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

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.

Upvotes: 6

Views: 140

Answers (1)

Ehsan
Ehsan

Reputation: 12397

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))

Upvotes: 3

Related Questions