如何屏蔽索引小于某个特定值的数组

Question

我有一个3D数组，其中第一个索引指的是高度。我有一个2D数组，其中每个元素都是最小高度

import numpy as np

a = np.ones((3,3,3)) # 3D array
b = [[1.2, 1.0, 2.0],
     [1.5, 1.2, 1.3],
     [1.0, 2.0, 0.5]]

我想屏蔽a，其中a的第一个索引/维度小于b给出的值。 例如：

a[0,1,1] = 0 and a[1,1,1] since b[1,1] = 1.2, but a[2,1,1] = 1

我的解决方案是使用for循环，但我想使用np.ma.mask（）创建一个布尔矩阵

我的解决方案：

nLat = a.shape[1]
nLon = a.shape[2]
for i in np.arange(0,nLat,1):
    for j in np.arange(0,nLon,1):
        minHeight = b[i,j]
        for hgt, value in enumerate(a):
            if hgt < minHeight:
                a[hgt,i,j] = 0

这将修改原始数组。虽然这是可行的，但我更愿意创建一个布尔数组（最好循环更少），然后将布尔值与原始值相乘，以创建一个最终输出，除了索引太小之外，该输出保持不变

Answer 1

我们可以通过与b的范围比较得到所需的掩码-

mask = np.arange(a.shape[0])[:,None,None]<b
a[mask] = 0

我们还可以使用内置进行外部比较，以获得遮罩：

mask = np.less.outer(np.arange(a.shape[0]),b)

如果您只对a的掩码等效项感兴趣，请使用-

L=3 # output length
a_mask = (np.arange(L)[:,None,None]>=b)