如何为二维 NumPy 数组中的 nd 值创建掩码？

Question

如果我想根据二维数组中的一维值创建掩码：

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

我得到：

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

如何为具有 nd 值的二维数组创建这样的掩码？就像下面二维数组中的二维值和二维阈值的例子：

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

看起来像这样：

[[[1 5]
  [3 5]]

 [[4 4]
  [7 2]]]

[1, 5]、[3, 5]、[4, 4] 和 [7, 2] 是示例性的 2d 值。在 threshold 中设置的阈值，第一个值为 2，第二个值为 4:

• [1, 5] 的单元格应为 False，因为 1 > 2 == False 和 5 > 4 == True
• [3, 5] 的单元格应为 True，因为 3 > 2 == True 和 5 > 4 == True
• [4, 4] 的单元格应为 False，因为 4 > 2 == True 和 4 > 4 == False
• [7, 2] 的单元格应为 False，因为 7 > 2 == True 和 2 > 4 == False

我需要做什么才能得到这个对应的面具？

[[ False  True]
 [ False False]]

Answer 1

numpy 广播比较实际上可以很好地为您处理这个问题。只需将 threshold 设为 1D 数组并沿最终轴断言 all。

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

(b > t).all(-1)

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

---

不过要澄清一下，您的数组实际上是 3D。如果您的数组是 2D，如下所示，这会有点不同：

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

(arr > t).all(-1)

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