finefoot
finefoot

Reputation: 11224

How to create a mask for nd values in a 2d NumPy array?

If I want to create a mask depending on a 1d value in a 2d array:

a = np.array([[3, 5], [7, 1]])
threshold = 2
mask = a > threshold
print(a)
print(mask)

I get:

[[3 5]
 [7 2]]
[[ True  True]
 [ True False]]

How can I create such a mask for a 2d array with nd values? Like the following example of 2d values and a 2d threshold in a 2d array:

b = np.array([[[1, 5], [3, 5]], [[4, 4], [7, 2]]])
threshold = 2, 4
print(b)

Looks like this:

[[[1 5]
  [3 5]]

 [[4 4]
  [7 2]]]

[1, 5], [3, 5], [4, 4] and [7, 2] are the exemplary 2d values. The threshold, as set in threshold, for the first value is 2 and for the second value it's 4:

What do I have to do to get this corresponding mask?

[[ False  True]
 [ False False]]

Upvotes: 0

Views: 580

Answers (1)

user3483203
user3483203

Reputation: 51165

numpy broadcasted comparison actually handles this quite nicely for you. Just make your threshold a 1D array and assert all along the final axis.

t = np.array([2, 4])

(b > t).all(-1)

array([[False,  True],
       [False, False]])

To clarify however, your array is actually 3D. If your array was 2D, like below, this would work a bit differently:

arr = np.array([[1, 5],
                [3, 5],
                [4, 4],
                [7, 2]])

(arr > t).all(-1)

array([False,  True, False, False])

Upvotes: 2

Related Questions