Iterating a function over an array

Question

Posing the title of this question differently.

I have a function that take a three dimensional array and masks certain elements within the array based on specific conditions. See below:

#function for array masking 
def masc(arr,z):
        return(np.ma.masked_where((arr[:,:,2] <= z+0.05)*(arr[:,:,2] >= z-0.05), arr[:,:,2])) 
    

arr is a 3D array and z is a single value.

I now want to iterate this for multiple Z values. Here is an example with 2 z values:

masked_array1_1 = masc(xyz,z1)
masked_array1_2 = masc(xyz,z2)

masked_1 = masked_array1_1.mask + masked_array1_2.mask
masked_array1 = np.ma.array(xyz[:,:,2],mask=masked_1)

The masked_array1 gives me exactly what i'm looking for.

I've started to write a forloop to iterate this over a 1D array of Z values:

mask_1 = xyz[:,:,2]
for i in range(Z_all_dim):
    mask_1 += (masc(xyz,Z_all[i]).mask)

masked_array1 = np.ma.array(xyz[:,:,2], mask = mask_1)

Z_all is an array of 7 unique z values. This code does not work (the entire array ends up masked) but i feel like i'm very close. Does anyone see if i'm doing something wrong?

Answer 1

Your issue is that before the loop you start with mask_1 = xyz[:,:,2] . Adding a boolean array to a float will cast the boolean to 1s and 0s and unless your float array has any 0s in it, the final array will be all nonzero values, which then causes every value to get masked. Instead do

mask_1 = masc(xyz, Z_all[0]).mask
for z in Z_all[1:]:
    mask_1 += masc(xyz, z).mask

Or avoiding any loops and broadcasting your operations

# No need to pass it through `np.ma.masked_where` if
# you're just going to extract just the boolean mask
mask = (xyz[...,2,None] <= Z_all + 0.05) * (xyz[...,2,None] >= Z_all - 0.05)
mask = np.any(mask, axis=-1)