Aggregate Numpy Functions

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 funct

Solution 1:

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

Solution 2:

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.