I have a 3D array where the first index refers to the height. I have a 2D array where each element is a minimum height.
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]]
I want to mask a
where the first index/dimension of a
is less than the value given by b
.
For example:
a[0,1,1] = 0 and a[1,1,1] since b[1,1] = 1.2, but a[2,1,1] = 1
My solution is to use for loops, but I would like to create a boolean matrix using np.ma.mask().
My solution:
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
This modifies the original array. While this works, I'd rather create a boolean array (preferably with fewer loops), and then multiply the boolean by the original to create a final output that is unchanged except where the indices are too small.
We can get the required mask with a ranged comparison with b
-
mask = np.arange(a.shape[0])[:,None,None]<b
a[mask] = 0
We can also use builtin for the outer comparison to get the mask :
mask = np.less.outer(np.arange(a.shape[0]),b)
If you are only interested in the mask equivalent of a
, use -
L=3 # output length
a_mask = (np.arange(L)[:,None,None]>=b)