Search code examples
pythonarraysnumpymasking

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


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.


Solution

  • 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)