jonalm jonalm - 3 months ago 18
Python Question

Index confusion in numpy arrays

I'm really confused by the index logic of numpy arrays with several dimensions. Here is an example:

import numpy as np
A = np.arange(18).reshape(3,2,3)
[[[ 0, 1, 2],
[ 3, 4, 5]],

[[ 6, 7, 8],
[ 9, 10, 11]],

[[12, 13, 14],
[15, 16, 17]]])


this gives me an array of shape (3,2,3), call them (x,y,z) for sake of argument. Now I want an array B with the elements from A corresponding to x = 0,2 y =0,1 and z = 1,2. Like

array([[[ 1, 2],
[4, 5]],

[[13, 14],
[16, 17]]])


Naively I thought that

B=A[[0,2],[0,1],[1,2]]


would do the job. But it gives

array([ 2, 104])


and does not work.

A[[0,2],:,:][:,:,[1,2]]


does the job. But I still wonder whats wrong with my first try. And what is the best way to do what I want to do?

Answer

There are two types of indexing in NumPy basic and advanced. Basic indexing uses tuples of slices for indexing, and does not copy the array, but rather creates a view with adjusted strides. Advanced indexing in contrast also uses lists or arrays of indices and copies the array.

Your first attempt

B = A[[0, 2], [0, 1], [1, 2]]

uses advanced indexing. In advanced indexing, all index lists are first broadcasted to the same shape, and this shape is used for the output array. In this case, they already have the same shape, so the broadcasting does not do anything. The output array will also have this shape of two entries. The first entry of the output array is obtained by using all first indices of the three lists, and the second by using all second indices:

B = numpy.array([A[0, 0, 1], A[2, 1, 2]])

Your second approach

B = A[[0,2],:,:][:,:,[1,2]]

does work, but it is inefficient. It uses advanced indexing twice, so your data will be copied twice.

To get what you actually want with advanced indexing, you can use

A[np.ix_([0,2],[0,1],[1,2])]

as pointed out by nikow. This will copy the data only once.

In your example, you can get away without copying the data at all, using only basic indexing:

B = A[::2, :, 1:2]
Comments