Python informer

Improve your Python coding skills

Avoiding loops in numpy

Suppose we had a numpy array a, and we wanted to add 1 to each element of the array. We could do it like this:

for i, x in enumerate(a):
    a[i] = x + 1

The problem here is that the for loop is being executed in Python. Now Python is fairly efficient, but it will usually be quite a bit slower than native code. And if the array is very large, that can be quite a performance hit.

Fortunately, numpy has a solution. You can just do this:

a = a + 1

numpy will then loop over the entire array, adding 1 to each element, and it will all be done in optimised, native code. Not only does this make the code easier to read, but it also makes it run faster.

This is just the simplest of many features in numpy that allow you to avoid explicit looping in most common cases. We will list the main ones here, with examples

Element-wise operations

As we saw above, numpy can apply an operator to all the elements in an array, with the need for a loop. You can actually create more complex expressions, like this:

import numpy as np

a = np.array([1, 2, 3, 4])
b = np.array([2, 4, 6, 8])
c = np.array([1, 4, 9, 16])
x = a*b + 5                #[ 7 13 23 37]

x is calculated by performing an element-wise calculation:

  • x[0] = a[0]*b[0] + 5 = 7
  • x[1] = a[1]*b[1] + 5 = 13
  • etc

In a similar way, y is calculated element-wise as a + c modulo 7:

y = (a + c) % 7            #[2 6 5 6]

In this case z is calculated using a comparison operator, so the result array contains boolean values:

z = b > c                  #[ True False False False]

In each case, the arrays must be of compatible shape. Usually this means that they need to be the same shape (in this case, a 1 dimensional array of 4 elements), but there are broadcast rules that allow for other situations.

Universal functions

Numpy also comes with a full set of mathematical functions that operate element-wise on arrays. For example, the square root function:

import numpy as np

a = np.array([1, 2, 4, 10])
x = np.sqrt(a)

Giving:

[1.         1.41421356 2.         3.16227766]