Removing NaN rows from a three dimensional array

Question

How can I remove the NaN rows from the array below using indices (since I will need to remove the same rows from a different array.

array([[[nan,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  0.],
        [ 0., nan,  0.,  0.],
        [ 0.,  0.,  0.,  0.]]])

I get the indices of the rows to be removed by using the command

a[np.isnan(a).any(axis=2)]

But using what I would normally use on a 2D array does not produce the desired result, losing the array structure.

a[~np.isnan(a).any(axis=2)]
array([[0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]])

How can I remove the rows I want using the indices obtained from my first command?

Answer 1

You need to reshape:

a[~np.isnan(a).any(axis=2)].reshape(a.shape[0], -1, a.shape[2])

But be aware that the number of NaN-rows at each 2D subarray should be the same to get a new 3D array.