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 (e.g. ([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?

hpaulj · Accepted Answer · 2017-01-19 19:33:03Z

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]]])

edited Jan 19, 2017 at 19:33

answered Jan 19, 2017 at 19:27

hpaulj

233k14 gold badges260 silver badges392 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Electronic Ant Over a year ago

I understand now. Thanks for the clear and concise answer!

Collectives™ on Stack Overflow

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

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related