user27118 user27118 - 1 month ago 16
Python Question

Efficiently summing outer product for 1D NumPy arrays

I have a function of the form

enter image description here

One way to implement this function in numpy is to assemble a matrix to sum over:

y = a*b - np.sum(np.outer(a*b, b), axis=0)


Is there a better way to implement this function with numpy, one that doesn't involve creating an NxN array?

Answer

You could use np.einsum -

y = a*b - np.einsum('i,i,j->j',a,b,b)

We can also perform a*b and feed to einsum -

y = a*b - np.einsum('i,j->j',a*b,b)

On the second approach, we can save some runtime by storing a*b and reusing.

Runtime test -

In [253]: a = np.random.rand(4000)

In [254]: b = np.random.rand(4000)

In [255]: %timeit np.sum(np.outer(a*b, b), axis=0)
10 loops, best of 3: 105 ms per loop

In [256]: %timeit np.einsum('i,i,j->j',a,b,b)
10 loops, best of 3: 24.2 ms per loop

In [257]: %timeit np.einsum('i,j->j',a*b,b)
10 loops, best of 3: 21.9 ms per loop