Menu
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.
935
The einsum function is one of NumPy's jewels. It can often outperform familiar array functions in terms of speed and memory efficiency, thanks to its expressive power and smart loops.
Then there’s a good chance einsum will help us do this much faster and more memory-efficiently that combinations of the NumPy functions multiply, sum and transpose would allow. As a small example of the function’s power, here are two arrays that we want to multiply element-wise and then sum along axis 1 (the rows of the array):
Additionally, np.einsum ('ij,jk', a, b) returns a matrix multiplication, while, np.einsum ('ij,jh', a, b) returns the transpose of the multiplication since subscript ‘h’ precedes subscript ‘i’. In explicit mode the output can be directly controlled by specifying output subscript labels.
You’d be forgiven for thinking that for a 3D array, np.einsum ('kij', M) moves the last axis to the first position and shifts the first two axes along. Actually, einsum creates its own output labelling by rearranging labels in alphabetical order. Therefore 'kij' becomes 'kij->ijk' and we have a sort of inverse permutation instead.
The tensordot function is also worth comparing for speed. If you search around, you’ll find examples of posts highlighting cases where einsum appears to be slow, especially when operating on several input arrays (such as this GitHub issue ). The einsum function was written by Mark Wiebe.
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))
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With
