Non-loopy way to create a dynamic mask-CodePudding

I am creating a mask based on the values in another array. Is there a non-loopy way to do so?

max_value = 5
num = 10
values = np.random.randint(max_value, size=(num))
mask = np.ones((num, max_value, max_value))
for idx in range(num):
    n = values[idx]
    mask[idx, -n:, -n:] = 0

CodePudding user response：

Just figured a non-loopy way - (a) create indices and broadcast them across the first two axes, (b) use inequality operation to create zeros and ones, (c) add with it's own transpose to get non-zero values at the masked positions. (caveat - if low is 0, it will not work).

max_value = 5
num = 10
values = np.random.randint(low=1, high=max_value, size=(num))
mask = np.ones((num, max_value, max_value), dtype=np.int64)
for idx in range(num):
    n = values[idx]
    mask[idx, -n:, -n:] = 0

# non-loopy 
indices = np.arange(max_value)[None, :].repeat(max_value, 0)
indices = indices[None, :].repeat(num, 0)
partial_mask = np.where(indices < max_value - values[:, None, None], 1, 0)
new_mask = np.where(partial_mask.swapaxes(-1, -2)   partial_mask, 1, 0)

print(np.all(mask == new_mask))

CodePudding user response：

You can use this code (mask is bool dtype, may by best use sparse matrices..):

mask = np.zeros((num, max_value, max_value), dtype=bool)
mask[np.arange(num),-values,-values]=True
mask=np.cumsum(mask, axis=2, dtype=bool)
mask=~np.cumsum(mask, axis=1,dtype=bool)