3

I have a numpy operation that I call intensively and I need to optimise:

 np.sum(a**2, axis=1)**.5   # where a is a 2 dimensional ndarray

This operation is composed of three functions and requires iterating through 'a' three times. It would be more efficient to aggregate all operations under one function and apply that function just once along axis 1. Unfortunately, numpy's apply_along_axis function is not an option as performance is around x1000 worse.

Is there a way of aggregating several numpy operations so it only has to loop once over the array?

Improve this question

2 Answers 2

Reset to default
4

When working with floating-point array, you can use np.einsum -

np.sqrt(np.einsum('ij,ij->i',a,a))

Runtime test -

In [34]: a = np.random.rand(1000,1000)

In [35]: np.allclose(np.sum(a**2, axis=1)**.5,np.sqrt(np.einsum('ij,ij->i',a,a)))
Out[35]: True

In [36]: %timeit np.sum(a**2, axis=1)**.5
100 loops, best of 3: 7.57 ms per loop

In [37]: %timeit np.sqrt(np.einsum('ij,ij->i',a,a))
1000 loops, best of 3: 1.52 ms per loop
Improve this answer
Sign up to request clarification or add additional context in comments.

1 Comment

Thank you. This is even better than I expected!
3

Take a look at numexpr, which allows you to evaluate numerical expressions faster than pure numpy:

In [19]: a = np.arange(1e6).reshape(1000,1000)

In [20]: import numexpr as ne

In [21]: %timeit np.sum(a**2,axis=1)**0.5
100 loops, best of 3: 6.08 ms per loop

In [22]: %timeit ne.evaluate("sum(a**2,axis=1)")**0.5
100 loops, best of 3: 4.27 ms per loop

The **0.5 is not part of the expression because the sum is a reduction operations and needs to be computed last in an expression. You could also run another evaluation for the sqrt/**0.5.

Improve this answer

2 Comments

numexpr as sqrt too, so maybe that could help to replace **0.5? Though that would mean nesting two evaluate calls probably.
The sqrtonly took only about 1% of the time of the sum, so I thought it was a bit of needless optimization at that point.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.