Select non-infinite data from 3D array in numpy

I try to filter columns with non-infinite data in the first row from a 3D array.

Here’s my try:

r1 is only an example how to select the entries that are larger than 2 in the first row.

r2: A, what I think, similar construction to filter non-infinite doesn’t work.
What would be the correct approach?

import numpy as np

data = np.array([
    [[1,2,3,4],      [5,6,7,8]],
    [[1,3,5,6],      [8,1,3,2]],
    [[np.Inf,1,1,8], [5,8,1,9]]
    ])

r1 = data[np.where(data[...,0] > 2)]
print(r1)

r2 = data[np.where(not np.isinf(data[...,0]))]
print(r2)

This gives the following. The result r1 is correct as it selects based on numbers > 2.

[[ 5.  6.  7.  8.]
 [ 8.  1.  3.  2.]
 [inf  1.  1.  8.]
 [ 5.  8.  1.  9.]]
Traceback (most recent call last): ...
r2 = data[np.where(not np.isinf(data[...,0]))]

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

>Solution :

You should use ~ (binary not), not not to invert the boolean array as ~ is broadcasting the negation to the elements while not would attempt to get a boolean value of the whole array (which is not supported as it is ambiguous):

r2 = data[~np.isinf(data[...,0])

output:

array([[1., 2., 3., 4.],
       [5., 6., 7., 8.],
       [1., 3., 5., 6.],
       [8., 1., 3., 2.],
       [5., 8., 1., 9.]])