numpy slice strange behavior

Viewed 105

I have a 5 dimension array like this

a=np.random.randint(10,size=[2,3,4,5,600])
a.shape #(2,3,4,5,600)

I want to get the first element of the 2nd dimension, and several elements of the last dimension

b=a[:,0,:,:,[1,3,5,30,17,24,30,100,120]]
b.shape #(9,2,4,5)

as you can see, the last dimension was automatically converted to the first dimension. why? and how to avoid that?

1 Answers

This behavior is described in the numpy documentation. In the expression

a[:,0,:,:,[1,3,5,30,17,24,30,100,120]]

both 0 and [1,3,5,30,17,24,30,100,120] are advanced indexes, separated by slices. As the documentation explains, in such case dimensions coming from advanced indexes will be first in the resulting array.

If we replace 0 by the slice 0:1 it will change this situation (since it will leave only one advanced index), and then the order of dimensions will be preserved. Thus one way to fix this issue is to use the 0:1 slice and then squeeze the appropriate axis:

a[:,0:1,:,:,[1,3,5,30,17,24,30,100,120]].squeeze(axis=1)

Alternatively, one can keep both advanced indexes, and then rearrange axes:

np.moveaxis(a[:,0,:,:,[1,3,5,30,17,24,30,100,120]], 0, -1)
Related