Index confusion in numpy arrays

Question

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 1

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]

Answer 2

I recommend the following advanced tutorial, which explains the various indexing methods: NumPy MedKit

Once you understand the powerful ways to index arrays (and how they can be combined) it will make sense. If your first try was valid then this would collide with some of the other indexing techniques (reducing your options in other use cases).

In your example you can exploit that the third index covers a continuous range:

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

You could also use

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

which is handy in more general cases, when the latter indices are not continuous. np.ix_ simply constructs three index arrays.

As Sven pointed out in his answer, there is a more efficient way in this specific case (using a view instead of a copied version).

Edit: As pointed out by Sven my answer contained some errors, which I have removed. I still think that his answer is better, but unfortunately I can't delete mine now.