numpy - Multidimensional boolean mask

Question

I'm quite new to Python and numpy and I just cannot get this to work without manual iteration.

I have an n-dimensional data array with floating point values and an equally shaped boolean "mask" array. From that I need to get a new array in the same shape as the both others with all values from the data array where the mask array at the same position is True. Everything else should be 0.:

# given
data = np.array([[1., 2.], [3., 4.]])
mask = np.array([[True, False], [False, True]])

# target
[[1., 0.], [0., 4.]]

Seems like numpy.where() might offer this but I could not get it to work.

Bonus: Don't create new array but replace data values in-position where mask is False to prevent new memory allocation.

Thanks!

Answer 1

This should work

data[~mask] = 0

Numpy boolean array can be used as index (https://docs.scipy.org/doc/numpy-1.15.0/user/basics.indexing.html#boolean-or-mask-index-arrays). The operation will be applied only on pixels with the value "True". Here you first need to invert your mask so False becomes True. You need the inversion because you want to operate on pixels with a False value.

Answer 2

Also, you can just multiply them. Because 'True' and 'False' is treated as '1' and '0' respectively when a boolean array is input in mathematical operations. So,

#element-wise multiplication
data*mask

or

np.multiply(data, mask)