python - How to quickly mask different slices in my array?

Question

I have a 3d array where all axis lengths are the same (for example (5,5,5)). I need to mask all of the array and keep certain slices in the array unmasked as per the code below. I managed to accomplish this using a for loop but I wondered if there was a faster solution out there.

array = np.reshape(np.array(np.random.rand(125)), (5,5,5))
array = ma.array(array, mask=True)
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        array[i, j, :].mask[i:j] = False

This allows me to sum this array with another array of the same size while ignoring the masked values.

标签： pythonnumpy

Answer 1

You can create the entire mask in one step using broadcasting:

i, j, k = np.ogrid[:5, :5, :5] 
mask = (i>k) | (k>=j)