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

How to mask an array where the index is less than a certain

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

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 的第一个 index/dimension 小于 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

这修改了原始数组。虽然这可行,但我宁愿创建一个布尔数组(最好有更少的循环),然后将布尔值乘以原始值以创建最终输出,除了索引太小的地方外,该输出不变。

我们可以通过 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)