The most useful Numpy function (if you work with tensors) 2013-02-06

I discovered np.einsum, the most useful of all Numpy library functions.

Suppose you want to compute:

$$ T^{svi} = \sum_x \sum_j P^{sx} M^{v}_{xj} Q^{ij} $$

Then you can just describe this operation using indices:

T = np.einsum("sx,vxj,ij -> svi", P, M, Q)

It's more compact than the LaTeX expression!

Another example:

$$ M^{v}_{xj} = \sum_{x} \sum_{j} T^{svi} P^{-1}_{sx} Q^{-1}_{ij} $$

becomes:

    M = np.einsum("svi,sx,ij -> vxj", T, P_inv, Q_inv)

One thought on “The most useful Numpy function (if you work with tensors)”

John Montgomery on October 30, 2013 at 4:44 am said:

Thanks for these examples, they were very helpful to me when I was figuring out what einsum does. I think there is a typo in the second one, the indices summed over are i and s, but the displayed equation shows x and j as the indices summed.

Comments are closed.