Search code examples
arraysnumpymask

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:

  • cell for [1, 5] should be False since 1 > 2 == False and 5 > 4 == True
  • cell for [3, 5] should be True since 3 > 2 == True and 5 > 4 == True
  • cell for [4, 4] should be False since 4 > 2 == True and 4 > 4 == False
  • cell for [7, 2] should be False since 7 > 2 == True and 2 > 4 == False

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

[[ False  True]
 [ False False]]

Solution

  • 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])