pythonnumpyslice

Unexpected behavior with array slicing and mask


It was unexpected that

x=np.empty((2,10,5))
x.shape
>>> (2, 10, 5)

x[0].shape, x[0,:,:].shape
>>> ((10, 5), (10, 5))

mask = [True,True,True,False,False]
x[0,:,mask].shape
>>> (3, 10)

I would have expected x[0,:,mask].shape to be (10,3) because the shape of x is (2,10,5) and applying the slice/mask this change was expected: (2→1, 10→10, 5→3) = (1,10,3).

It seems however transposing the result?

For another array it works as expected:

y=np.empty((2,5))
y.shape
>>> (2, 5)

y[0].shape, y[0,:].shape
>>> ((5,), (5,))

y[:,mask].shape
>>> (2, 3)

Solution

  • NumPy boolean masks don't just select elements but restructure the output array. When you use a mask on a dimension, it collapses that dimension and moves the selected elements to the first axis of the result, causing the unexpected transposition in your multidimensional array.

    If you want to preserve the original dimensions, you can use numpy.compress(condition, a, axis=None, out=None)

    result = np.compress(mask, x[0], axis=1) 
    print(result.shape)
    
    output
    (10, 3)