NumPy indexing explained. NumPy is the universal standard for... | by Àlex Escolà Nixon | Towards Data Science

NumPy is the universal standard for working with Numerical data in Python. Multidimensional NumPy arrays are extensively used in Pandas, SciPy, Scikit-Learn, scikit-image, which are some of the main data science and scientific Python packages. So a good understanding of NumPy is crucial if we are working with these tools! And so it is to understand its indexing methods, which we will go through in this post.
Note that a key aspect or motivation in using NumPy, is that it enables to manipulate large quantities of numerical data in a vectorised way, that is, avoiding the use of inefficient python loops. The huge advantage of vectorising your code is that you’re pushing all looping down to C level, which is much faster.
Now let’s get to the topic of the post, which is indexing NumPy arrays. Indexing with NumPy can be a little tough and counter-intuitive without a good understanding of its basic rules. Here we’ll dive into the different indexing methods that can be used, which are:
- Basic Slicing and Indexing
- Advanced Indexing (including both integer and boolean array indexing)
Find a notebook version of this post along with all explanations here.
NumPy’s basic slicing is an extension of Python’s basic slicing concept extended to N dimensions.
It essentially allows you to take slices of an array along its dimensions using basic slicing notation, i.e start:stop:step. If you are not too familiar with Python’s basic slicing notation you can check this post in stack overflow which makes it very clear. Let’s see with an example:
array([['A', 'B', 'C', 'D', 'E'],       ['F', 'G', 'H', 'I', 'J'],       ['K', 'L', 'M', 'N', 'O'],       ['P', 'Q', 'R', 'S', 'T'],       ['U', 'V', 'W', 'X', 'Y']])
Yes, I know z is missing... Nothing against it, but we need a homogeneous matrix. This example array has a shape of (5, 5), so two dimensions. We can take slices along its dimensions (I’ll be using the equivalent term axes) by following the mentioned notation, for instance:
a[:2, 1:4]array([['B', 'C', 'D'],       ['G', 'H', 'I']])
So far so good. We’ve basically sliced the array along the first axis to get up to the second row starting at index 0 (note that the stop index is not included!), and along the second to get columns one to three. Think of the axes just as the (x,y,z) dimensions of a matrix. Where a third dimension can be simply though of as stacking multiple 2D arrays together:
Okay, so now we know we can slice an n-dimensional array using python’s slice notation. But what if you had to take values from a given axis using a list of indexes?
Well, here’s where advanced indexing comes into play ↓
Let’s just go with an example of what an intuitive way of doing this would be:
a_indexed = a[[4,3,1], [2,4,0]]
Okay, so following the same logic as in basic slicing, one could expect to get all values that fall into the above indexes along the different axes of the ndarray. Namely all values in rows 4, 3 and 1 which are in columns 2, 4 and 0 respectively, so the intersecting values that are highlighted bellow:
Is that so? Let’s check:
print(a_indexed)array(['W', 'T', 'F'])
Clearly not... (example inspired by this very instructive talk by Jaime Fernández :) ) What happened there? Why are we getting a one-dimensional result?
This is because advanced indexing follows a different set of rules. A good way to think about it, is that when using basic slicing, we are indexing on a grid which is defined by the slices we take on each dimension. Whereas using advanced indexing can be thought of as specifying a set of (x,y) coordinates of the values we want to retrieve.
In the above case indexing with [4,3,1],[2,4,0] , will behave as indexing on(4,2) , (3,4) and (1,0) respectively. So each index we specify in each dimension, will be combined with the corresponding indexes from the other dimensions.
More generally, when using advanced indexing, we must take into account these two main aspects (from the docs):
The indexing arrays represent a number of indexes into that dimension
Which basically means that we will retrieve as many elements from a given axis as indexes specified in the indexing array into that dimension. In the above case, we’ve retrieved 3 elements specified by the length 3 indexing arrays.
The result shape is identical to the (broadcast) indexing array shapes
To understand the above we need to dive a little into broadcasting...
Going deep into broadcasting could easily lead to an another entire blog post, so i’ll just try to cover the very essentials... From the docs:
broadcasting is defined as a term describing how NumPy treats arrays with different shapes during arithmetic operations. Subject to certain constraints, the smaller array is “broadcast” across the larger array so that they have compatible shapes
So basically, when some operation involving arrays with different shapes is performed, NumPy tries to make their shapes compatible before the operation takes place. Let’s take a look at some visual examples:
In the first case for instance, a one-dimensional NumPy array is added to an integer. So under the hood, before the addition takes place NumPy will broadcast the smaller shaped array across the larger, i.e it will replicate its value until its shape becomes compatible with that of the larger array.
And how do we know if two shapes are compatible?
We can find the general broadcasting rules defined in the docs. It is stated that two dimensions are compatible when
they are equal, or
one of them is 1
Okay so what does all this have to do with advanced indexing? Well, reminding the second mentioned point, the resulting shape from an advanced indexing is identical to the (broadcast) indexing array shapes. This basically means that NumPy will try to make the shapes from the indexing arrays compatible before performing the indexing operation. A simple way to inspect what the resulting shape will look like (in the case the arrays can be broadcast) is by using np.broadcast . Let’s see with the same example as above:
a = array([['A', 'B', 'C', 'D', 'E'],           ['F', 'G', 'H', 'I', 'J'],           ['K', 'L', 'M', 'N', 'O'],           ['P', 'Q', 'R', 'S', 'T'],           ['U', 'V', 'W', 'X', 'Y']])
Say we now want to index the array with:
rows = np.array([0,2,1])cols = np.array([2])
So using the above indexes, we would get an array with the shape:
ix = np.broadcast(rows, cols)print(ix.shape)(3,)
So here as we can see, a broadcast as in the first case in the figure above is taking place. The array will be indexed on the following (x,y) coordinates:
print(*ix)(0, 2), (2, 2), (1, 2)
Hence we get:
a[rows, cols]array(['C', 'M', 'H'])
Which, since both arrays are being broadcast, the above is equivalent to:
a[[0, 2, 1], [2, 2, 2]]array(['C', 'M', 'H'])
So in order to index with integer arrays, we just have to make sure that their shapes can be broadcast together and then just let NumPy take care of the rest.
Okay, so now we know how to index an array using integer indexes...
So the question that follows is....
What if I want the indexing to behave as it does with basic slicing? i.e how can we obtain what we were expecting in the first case, instead of array(['W', 'T', 'F']) ?
Let’s go back to that example... We were using the indexes:
rows = np.array([4,3,1])cols = np.array([2,4,0])
In this case, as we verified before, we are getting a one-dimensional array as a result from:
ix = np.broadcast(rows, cols)print(*ix)(4, 2) (3, 4) (1, 0)
So we have to find a way for NumPy to understand that we want to retrieve a grid containing all values. How can this be done...?
The answer is... to USE BROADCASTING!!
Taking into account that what broadcasting is doing is to stretch the smaller array so that all array shapes are compatible, we can leverage this idea to get the result we expected.
What we can do is to add an axis to one of the arrays, so that NumPy broadcasts the smaller to fit the size of the larger (smaller here in terms of the amount of dimensions). That can be done with:
rows[:,np.newaxis]array([[4],       [3],       [1]])
Or we could equivalently use rows[:,None]. Now both indexing arrays con be broadcast together and can be used to index the array:
ix = np.broadcast(rows[:,None], cols)print(*ix)(4, 2) (4, 4) (4, 0) (3, 2) (3, 4) (3, 0) (1, 2) (1, 4) (1, 0)
We can think of this as generating the cartesian product of our lists of indexes in order to retrieve all of them from the array. We could actually reproduce this using itertools.product :
print(*product(rows, cols))(4, 2) (4, 4) (4, 0) (3, 2) (3, 4) (3, 0) (1, 2) (1, 4) (1, 0)
So now by indexing the array with the broadcast indexes, we get:
a[rows[:,None], cols]array([['W', 'Y', 'U'],       ['R', 'T', 'P'],       ['H', 'J', 'F']])
This broadcasting can also be achieved using the function np.ix_:
a[np.ix_(rows, cols)]array([['W', 'Y', 'U'],       ['R', 'T', 'P'],       ['H', 'J', 'F']])
So as we can see integer indexing is a really useful tool, we just need to understand how it works. A very common application is when we have an array or nested list of indexes, for instance:
cols = [[3,4], [0,2], [0,1], [1,2], [3,3]]
And we want to take these columns respectively along the first axis, so columns 3 and 4 from the first row, 0 and 2 from the second row and so on. How can we do that? Why don’t you give it a try...?
As mentioned earlier, in integer indexing indexes along the dimensions are combined through broadcasting rules. Hence, leveraging broadcasting, we want to obtain a set of indexes whose combined shape is broadcastable to the desired output shape. Which in this case is straight forward, since we have a sublist for each row. So it becomes clear that we want a range containing as many values as rows to index across the first axis.
So an obvious way to do this could be:
rows = np.arange(a.shape[0])a[rows, cols]
However, this would yield:
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (5,) (5,2)
As we can see the error is quite explicit, the indexing arrays cannot be broadcast together. This is because the shapes along the corresponding axis are not compatible, since we have that:
cols (2d array): 5 x 2rows (1d array):     5 
With the above sketch it becomes clear that for the shapes to be compatible, we must add a new axis to rows, so that the second axes on both arrays have the same shape and the first axis of rows is 1, which satisfies the second rule of broadcasting:
cols         (2d array): 5 x 2rows[:,None] (1d array): 5 x 1
Now we would get as expected:
a[rows[:,None], cols]array([['D', 'E'],       ['F', 'H'],       ['K', 'L'],       ['Q', 'R'],       ['X', 'X']])
Now for the last section... Is it possible to combine both types of indexing?
WARNING...! (this last section gets a little more complicated, and you’ve already gone through the most relevant and useful stuff, so feel free to take a break, or leave it here if you think you’ve had enough for now :)
It is also possible to combine basic slicing and advanced indexing. This will lead to integer array indexes being broadcast together in the same fashion as we’ve previously seen, and the slices behaving as we’ve seen in the basic slicing section. However under some situations this can lead to unexpected results.
Let’s take the following example as a reference to better explain this:
a = np.arange(60).reshape(3,4,5)print(a)array([[[ 0,  1,  2,  3,  4],        [ 5,  6,  7,  8,  9],        [10, 11, 12, 13, 14],        [15, 16, 17, 18, 19]],       [[20, 21, 22, 23, 24],        [25, 26, 27, 28, 29],        [30, 31, 32, 33, 34],        [35, 36, 37, 38, 39]],       [[40, 41, 42, 43, 44],        [45, 46, 47, 48, 49],        [50, 51, 52, 53, 54],        [55, 56, 57, 58, 59]]])
For combined indexing, we need to take into account these following guiding rules:
Integer array indexing rules apply only if two or more non-slice objects are introduced
That means that if we only use one integer array to index such as:
a_s = a[:2, [0, 1]]a_s.shape(2, 2, 5)print(a_s)array([[[ 0,  1,  2,  3,  4],        [ 5,  6,  7,  8,  9]],       [[20, 21, 22, 23, 24],        [25, 26, 27, 28, 29]]])
So since basic slicing rules apply here the above would be equivalent to a[:2,:2].
However, as soon as we have more than one indexing array, advanced indexing rules will apply:
a_s = a[:2, [0, 1], [2, 4]]a_s.shape(2, 2)print(a_s)array([[ 2,  9],       [22, 29]])
We we can see that adding an indexing array for the last axis, leads to an output the shape of which is dependant of the broadcast shape of both indexing arrays, which is (2,2).
The resulting axes introduced by the arrays indexes are at the front, unless they are consecutive
Let’s take a look at some examples to see this more clearly. Taking for instance the following pairs of indexes:
ix1 = np.array([0,1])ix2 = np.array([1,3])
We could for instance, index along the two first axes using ix1 and ix2 respectively and take a full slice along the last axis. In such case we would get the 2D-arrays 0 and 1, rows 1 and 3 respectively and a full slice along the last axis for both cases:
a_s = a[ix1, ix2, :]a_s.shape(2, 5)print(a_s)array([[ 5,  6,  7,  8,  9],       [35, 36, 37, 38, 39]])
So we’ve selected the elements:
Note that as mentioned earlier, here both different method’s rules work as expected and are combined. Basic slicing along the last axis is following the rules of basic slicing mentioned in the first section, so if you were to instead index with an array you’d get:
ix3 = np.array([1,2])a[ix1, ix2, ix3]array([ 6, 37])
We could also index in the following way:
a_s = a[:, ix1, ix2]a_s.shape(3, 2)print(a_s)array([[ 1,  8],       [21, 28],       [41, 48]])
So a full slice on the first axis to select all 2D arrays, then select rows 0 and 1 and columns 1 and 3 respectively, so:
So finally if we use the indexing arrays to index on the first and last axes, as a[ix1, :, ix2], we would expect to index on the first two 2D-arrays, take a full slice on the second axis (so all rows), and columns 1 and 3 respectively. So:
With a resulting shape of (4, 2), let’s check:
a_s = a[ix1, :, ix2]s_s.shape(2,4)print(a_s)array([[ 1,  6, 11, 16],       [23, 28, 33, 38]])
Why is this case different...??
Going back to the second rule: The resulting axes introduced by the arrays indexes are at the front, unless they are consecutive. So since the indexing arrays here are not consecutive, the resulting axes on which they’ve been used will come at the front, and the sliced dimension at the back.
While it might seem very weird when indexing with 1-dimensional arrays, take into account that it is also possible to index with arrays of an arbitrary amount of dimensions. Say we are indexing the same example array both on the first and last axes with 3d arrays, both of shape (3,4,2). So we know that the final array will somewhere also have the shape (3,4,2), since both indexing arrays broadcast to the same shape. Now the question is, where do we place a full slice taken between the first and last axis?
Given that it is no longer as clear that it should go in the middle, there is a convention in these cases which is that sliced dimensions go at the end.
So in such cases it will be our task to rearrange the dimensions of the array to match our expected output. On the example above what we could do is to swap the last two axes and get as we expected with:
a[ix1, :, ix2].swapaxes(0,1)array([[ 1, 23],       [ 6, 28],       [11, 33],       [16, 38]])
Which is basically the same as transposing in this case, a[ix1, :, ix2].T.
So, as a takeaway from this last section, we must be careful when combining both indexing methods since we might not always get what we expect. The safest will be to bear in mind the caveats of combined indexing to make sure that it will behave as we expect it to.
Thanks a lot for taking the time to read this post and I hope you enjoyed it :)