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

user_12919727 0
Viewed 290

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.

masking numpy arrays python
1 Answers

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)
user_32938810
Related