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?

Divakar
  • 218,885
  • 19
  • 262
  • 358
prl900
  • 4,029
  • 4
  • 33
  • 40

2 Answers2

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
Divakar
  • 218,885
  • 19
  • 262
  • 358
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.

Jan Christoph Terasa
  • 5,781
  • 24
  • 34
  • `numexpr` as `sqrt` too, so maybe that could help to replace `**0.5`? Though that would mean nesting two `evaluate` calls probably. – Divakar Mar 15 '16 at 07:47
  • The `sqrt`only took only about 1% of the time of the `sum`, so I thought it was a bit of needless optimization at that point. – Jan Christoph Terasa Mar 15 '16 at 09:00