6

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.

semmo
  • 63
  • 4

1 Answers1

3

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))
Ehsan
  • 12,072
  • 2
  • 20
  • 33
  • 1
    Good idea there. To get max perf, think you can skip `optimize` for `evec2`, as there's no sum reduction there. – Divakar Jul 09 '20 at 07:15
  • @Divakar, That is true. Thank you for the point. I edit the post. – Ehsan Jul 09 '20 at 07:33
  • 1
    Wow! It's blazing fast. It went from 16 seconds to 0.007 seconds. I'll make sure to try to reuse calculations from now on. Thanks! – semmo Jul 09 '20 at 08:10