Does anyone know if it's possible to use a forloop to create arrays and store them as unique variables?

For example. I have a function that takes a three-dimensional array and masks elements within the array based on given parameters:

def masc(arr,M,N,O,P,z):
        return(np.ma.masked_where((arr[:,:,2] <= z 0.05)*(arr[:,:,2] >= z-0.05)*(arr[:,:,0]>=M)*((arr[:,:,0])<=O)*(arr[:,:,1]<=N)*((arr[:,:,1])>=P-10), arr[:,:,2]))

where arr is a three dimensional array.

my array is defined by variable 'xyz' and is a grid x columns and y rows with corresponding z values. Something like this:

xyz = 
[[40,50,2],
[41,50,5]
.....
[57,65,15]]

(these are totally made up numbers)

I run this function iteratively to create unique masked arrays based on a xyz and given variables. The output is the array xyz with certain elements masked based on the subsetting given in the masc function.

masked_array1 = masc(xyz,x1,y1,X1,Y1,z1)
masked_array2 = masc(xyz,x2,y2,X2,Y2,z2)
.
.
.
masked_arrayn = masc(xyz,xn,yn,Xn,Yn,zn)

I'd like to put this in a forloop where I can generate all of those unique arrays based on a given input arrays. Say it would work through an array of each variable:

X = [1,2,3]
Y = [1,2,3]
x = [4,5,6]
y = [4,5,6]
z = [7,8,9]

and the for loop would calculate the corresponding masked array for each one of those values. So if I were to just do it manually myself it would look like:

masked_array1 = masc(xyz,x[1],y[1],X[1],Y[1],z[1])
masked_array2 = masc(xyz,x[2],y[2],X[2],Y[2],z[2])
masked_array3 = masc(xyz,x[3],y[3],X[3],Y[3],z[3])

Does anyone know if it's possible to use a forloop to create unique variables such that I'd achieve multiple output arrays? Would I have to create one large array that encompasses all the masked arrays?

EDIT: I went in and added some detail and example input output. Hopefully that helps clarify my question.

CodePudding user response：

For your function and some data

import numpy as np

def masc(arr,M,N,O,P,z):
    return(np.ma.masked_where((arr[:,:,2] <= z 0.05)*(arr[:,:,2] >= z-0.05)*(arr[:,:,0]>=M)*((arr[:,:,0])<=O)*(arr[:,:,1]<=N)*((arr[:,:,1])>=P-10), arr[:,:,2]))

arr = np.random.rand(10, 10, 3)

X = [1,2,3]
Y = [1,2,3]
x = [4,5,6]
y = [4,5,6]
z = [7,8,9]

you can iterate over all your parameters like

for xx, yy, XX, YY, zz in zip(x, y, X, Y, z):
    masc(arr, xx, yy, XX, YY, zz)

or create a single output array like

out1 = np.array([
    masc(arr, xx, yy, XX, YY, zz)
    for xx, yy, XX, YY, zz in zip(x, y, X, Y, z)
])

alternatively you can use broadcasting

def masc2(arr,M,N,O,P,z):
    M = np.asarray(M)[..., None, None]
    N = np.asarray(N)[..., None, None]
    O = np.asarray(O)[..., None, None]
    P = np.asarray(P)[..., None, None]
    z = np.asarray(z)[..., None, None]

    condition = (
        (arr[:,:,2] <= z   0.05) *
        (arr[:,:,2] >= z - 0.05) *
        (arr[:,:,0] >= M) *
        (arr[:,:,0] <= O) *
        (arr[:,:,1] <= N) *
        (arr[:,:,1] >= P - 10)
    )
    data = arr[:,:,2]
    return np.ma.masked_where(condition, np.broadcast_to(data, condition.shape))

and do the entire operation in one single call

out2 = masc2(arr, x, y, X, Y, z)

and the outputs are identical

np.allclose(out1, out2)
# True