Filter out records from 3D array based on the values in column of 2D array-CodePudding

I have a 3D array of shapes (3, 5, 5) as per below. I want to filter out all records (rows) from constituent 2D arrays where col of index 0 is nan.

I have tried below but it returns back a 2D array instead of a desired 3D array.

arr[ np.where(arr[:, :, 0] != np.nan)]

Can you please suggest the correct way of achieving the stated objective?

Thank you!

[[[1. 0. 0.10. 0.]
  [2. 0. 0. 9. 0.]
  [nan. 0. 0. 8. 0.]
  [4. 0. 0. 7. 0.]
  [nan. 0. 0. 6. 0.]]
 [[1. 0. 0. 199. 0.]
  [2. 0. 0. 198. 0.]
  [3. 0. 0. 196. 0.]
  [nan. 0. 0. 190. 0.]
  [nan. 0. 0. 160. 0.]]
 [[1. 0. 0. 999. 0.]
  [2. 0. 0. 870. 0.]
  [nan. 0. 0. 270. 0.]
  [nan. 0. 0. 100. 0.]
  [nan. 0. 0. 80. 0.]]]

CodePudding user response：

You can try to use a list of 2D arrays as a result because after na filtering the shapes of inner the 2D arrays will be different so they can not be stacked together into a 3D array.

import numpy as np

arr = np.array([
 [[1., 0., 0., 10., 0.],
  [2., 0., 0., 9., 0.],
  [np.nan, 0., 0., 8., 0.],
  [4., 0., 0., 7., 0.],
  [np.nan, 0., 0., 6., 0.]],
 [[1., 0., 0., 199., 0.],
  [2., 0., 0., 198., 0.],
  [3., 0., 0., 196., 0.],
  [np.nan, 0., 0., 190., 0.],
  [np.nan, 0., 0., 160., 0.]],
 [[1., 0., 0., 999., 0.],
  [2., 0., 0., 870., 0.],
  [np.nan, 0., 0., 270., 0.],
  [np.nan, 0., 0., 100., 0.],
  [np.nan, 0., 0., 80., 0.]]],
)

row_masks = np.isnan(arr[:, :, 0])
result = [arr2d[~row_mask] for arr2d, row_mask in zip(arr, row_masks)]

Result:

[array([[ 1.,  0.,  0., 10.,  0.],
        [ 2.,  0.,  0.,  9.,  0.],
        [ 4.,  0.,  0.,  7.,  0.]])
 array([[  1.,   0.,   0., 199.,   0.],
        [  2.,   0.,   0., 198.,   0.],
        [  3.,   0.,   0., 196.,   0.]])
 array([[  1.,   0.,   0., 999.,   0.],
        [  2.,   0.,   0., 870.,   0.]])]