Numpy: slicing a 3D array from a 4D array?

Question

Say if I have a 4D numpy array of dimensions: 10 x 10 x 10 x 10, and I want to obtain 3D subset based on some indexes I know (eg ([1,2],0,[5,6],[9,0]). How do I go about doing this, and what should I do if my indices aren't sequential, like in my example?

For instance, if I try doing this:

test[0:2, 0, 3:5, 0:2]

I get a 2 x 2 x 2 array as expected. However, if I try explicitly defining each dimension value like so:

test[[0,1], 0, [3,4], [0,1]]

I instead obtain a 2 x 1 array. What am I doing wrong?

Answer 1

The difference is between basic indexing and advanced. Indexing with slices and scalars is 'basic'. When one or more of the indexes are lists or array, it uses 'advanced'. Docs on advanced v basic indexing

There are various ways of producing the kind of indexing that you want. np.ix_ is designed to make this easy:

In [80]: np.ix_([0,1], [0], [3,4], [0,1])
Out[80]: 
(array([[[[0]]],


        [[[1]]]]), array([[[[0]]]]), array([[[[3],
          [4]]]]), array([[[[0, 1]]]]))

In [79]: arr[np.ix_([0,1], [0], [3,4], [0,1])].shape
Out[79]: (2, 1, 2, 2)

To pass all indices through ix_ I had to make the 2nd one a list. A way around that is to use ix_ on just 3 lists, and add that 0 later:

In [81]: I,J,K=np.ix_([0,1], [3,4], [0,1])
In [82]: arr[I,i,J,K].shape
Out[82]: (2, 2, 2)
In [83]: I,J,K
Out[83]: 
(array([[[0]],

        [[1]]]), array([[[3],
         [4]]]), array([[[0, 1]]]))

Note the shape of the index arrays:

In [84]: I.shape
Out[84]: (2, 1, 1)
In [85]: J.shape
Out[85]: (1, 2, 1)
In [86]: K.shape
Out[86]: (1, 1, 2)

together they broadcast to a (2,2,2) shape. So any index arrays that broadcast that way should work. You could make J with

In [87]: J
Out[87]: 
array([[[3],
        [4]]])
In [88]: np.array([3,4])[None,:,None]
Out[88]: 
array([[[3],
        [4]]])