2

I encountered a quite weird problem when indexing a numpy ndarray. You can produce the result with following code. I don't understand why the result of indexing a is somehow transposed while the result of 2d array b is normal. Thanks.

In [1]: a = np.array(range(6)).reshape((1,2,3))

In [2]: mask = np.array([True, True, True])

In [3]: a
Out[3]: 
array([[[0, 1, 2],
        [3, 4, 5]]])

In [4]: a[0, :, mask]
Out[4]: 
array([[0, 3],
       [1, 4],
       [2, 5]])

In [5]: a[0, :, mask].shape
Out[5]: (3, 2)

In [6]: b = np.array(range(6)).reshape((2,3))

In [7]: b[:, mask].shape
Out[7]: (2, 3)
Texot
  • 53
  • 5

1 Answers1

3

a[0, :, mask] mixes advanced indexing with slicing. The : is a "slice index", while the 0 (for this purpose) and mask are consider "advanced indexes".

The rules governing the behavior of indexing when both advanced indexing and slicing are combined state:

There are two parts to the indexing operation, the subspace defined by the basic indexing (excluding integers) and the subspace from the advanced indexing part. Two cases of index combination need to be distinguished:

  • The advanced indexes are separated by a slice, ellipsis or newaxis. For example x[arr1, :, arr2].
  • The advanced indexes are all next to each other. For example x[..., arr1, arr2, :] but not x[arr1, :, 1] since 1 is an advanced index in this regard.

In the first case, the dimensions resulting from the advanced indexing operation come first in the result array, and the subspace dimensions after that. In the second case, the dimensions from the advanced indexing operations are inserted into the result array at the same spot as they were in the initial array (the latter logic is what makes simple advanced indexing behave just like slicing).

So since a[0, :, mask] has advanced indexes separated by a slice (the first case), the shape of the resulting array has the axes associated to the advanced indexes pushed to the front and the axes associated with the slice pushed tho the end. Thus the shape is (3, 2) since the mask is associated with the axis of length 3, and the slice, :, associated with the axis of length 2. (The 0 index in effect removes the axis of length 1 from the resultant array so it plays no role in the resultant shape.)

In contrast, b[:, mask] has all the advanced indexes together (the second case). So the shape of the resulting array keeps the axes in place. b[:, mask].shape is thus (2, 3).

unutbu
  • 842,883
  • 184
  • 1,785
  • 1,677