CODEWITHSUNDEEP
pythonnumpyoptimization
user27118
user27118

Reputation: 147

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?

Upvotes: 1

Views: 62

Answers (1)

Divakar
Divakar

Reputation: 221584

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

Upvotes: 2

Related Questions