How can I simply filter a 3d numpy array by its 1st column values?

user_9712193 0
Viewed 2082

Suppose I have a 3D numpy array like this:

 data = np.array([[[1,2,3,4],[1,2.5,3,5]],
                 [[116,230,450,430],[80,100,300,320]],
                 [[60,100,120,80],[50,80,100,90]]])

How can I simply extract from it a 3D numpy array of same shape with a condition on axis 0, for example selecting those "rows" for which axis 0 < 3? A naïve way would be

data[data[0]<3]

But this fails:

IndexError: boolean index did not match indexed array along dimension 0; dimension is 3 but corresponding boolean dimension is 2

numpy arrays python
1 Answers

See my comment above, but from your data I am guessing you want the rows with any values less than 3. If so you could do:

data[(data<3).any(axis=2)]
>>> array([[1. , 2. , 3. , 4. ],
           [1. , 2.5, 3. , 5. ]])

EDIT1:

Solution can be achieved using transposition to match up the axis dimensions:

data.T[(data[0]<3).any(axis=0).T].T
>>> array([[[  1. ,   2. ],
            [  1. ,   2.5]],

           [[116. , 230. ],
            [ 80. , 100. ]],

           [[ 60. , 100. ],
            [ 50. ,  80. ]]])

EDIT2:

Another method that does not involve transposing. To apply the mask (data[0]<3).any(axis=0) onto the original data array the axes shapes must match. The shape of the mask is (4,) and data.shape = (3, 2, 4), so we need to apply the mask to the last axis as:

data[..., (data[0]<3).any(axis=0)]
>>> array([[[  1. ,   2. ],
            [  1. ,   2.5]],

           [[116. , 230. ],
            [ 80. , 100. ]],

           [[ 60. , 100. ],
            [ 50. ,  80. ]]])
user_120631260
Related