Search code examples
pythonarraysnumpymasking

create a 3d binary mask in numpy with a 2d array of depth values

I would like to create a 3d mask array based on a 2d array of 'depths' that indicate the extent to which the third dimension is filled with ones. I had hoped there would've been a vectorized solution to this within numpy rather than resorting to for loops. For example:

tst = np.arange(1, 10).reshape(3,3) # depths referred to above

blk = np.zeros(shape=(3,3,9)) # the output 3d array to be filled with ones or zeros

The intended results at the following example positions would be:

blk[0,0]
array([1., 0., 0., 0., 0., 0., 0., 0., 0.])
blk[1,2]
array([1., 1., 1., 1., 1., 1., 0., 0., 0.])
blk[2,2]
array([1., 1., 1., 1., 1., 1., 1., 1., 1.])

Does anyone have any ideas/hints on this?

Solution

  • You can use broadcasting, comparing to numpy.arange:

    blk = (tst[..., None] > np.arange(tst.max())[None, None]).astype(int)

    Output:

    array([[[1, 0, 0, 0, 0, 0, 0, 0, 0],     # blk[0, 0]
        [1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0]],

       [[1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0]],    # blk[1, 2]

       [[1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1]]])   # blk[2, 2]