CODEWITHSUNDEEP
pythonnumpyvectorizationmasked-array
Ho-Hin Leung
Ho-Hin Leung

Reputation: 1

How to apply a function element-wise with inputs from multiple numpy masked arrays to create a new masked array?

I have a function that takes in 4 single value inputs to return a singular float output, for example:

from scipy.stats import multivariate_normal

grid_step = 0.25 #in units of sigma
grid_x, grid_y = np.mgrid[-2:2+grid_step:grid_step, -2:2+grid_step:grid_step]
pos = np.dstack((grid_x, grid_y))
rv = multivariate_normal([0.0, 0.0], [[1.0, 0], [0, 1.0]])
grid_pdf = rv.pdf(pos)*grid_step**2
norm_pdf = np.sum(rv.pdf(pos))*grid_step**2

def cal_prob(x, x_err, y, y_err):
    x_grid = grid_x*x_err + x
    y_grid = grid_y*y_err + y
    PSB_grid = ((x_grid>3) & (y_grid<10) & (y_grid < 10**(0.23*x_grid-0.46)))
    PSB_prob = np.sum(PSB_grid*grid_pdf)/norm_pdf
    return PSB_prob

What this function is doing is estimating the probability that some x-y measurement is within some defined limit in x-y space, given x and y's uncertainties. It assumes the uncertainties are Gaussian and uncorrelated. Then, using the pre-made grid_pdf, it checks which grid points (scaled by x_err/y_err and shifted by x/y) are within the defined limit, and multiply the True/False grid by the grid_pdf, normalized by norm_pdf. The probability is given by the sum of the normalized array.

I want this function to be applied element-wise with those 4 inputs stored in 4 separate numpy masked arrays of the same shape, with possibly different masks, then use the function outputs to create a new array of the same shape. Is there a way that doesn't use a for loop?

Thanks!

My current solution is this:

mask1 = np.array([[False, True, False],[True, True, True],[True, False, False]])
mask2 = np.array([[True, True, True],[True, True, False],[False, False, True]])
# the only overlaps should be [0,1], [1,0] and [1,1]

x = np.ma.array(np.random.randn(*mask1.shape), mask=~mask1)
x_err = np.ma.array(np.abs(np.random.randn(*mask1.shape))*0.1, mask=~mask1)

y = np.ma.array(np.random.randn(*mask2.shape), mask=~mask2)
y_err = np.ma.array(np.abs(np.random.randn(*mask2.shape))*0.1, mask=~mask2)

# a combined mask to iterate through
all_mask = x+x_err+y+y_err

prob = np.zeros(mask1.shape)
prob = np.ma.masked_where(np.ma.getmask(all_mask), prob)

for i,xi in np.ma.ndenumerate(all_mask):
    prob[i] = cal_prob(xi, x_err[i], y[i], y_err[i])

Upvotes: 0

Views: 212

Answers (1)

hpaulj
hpaulj

Reputation: 231738

A test of np.vectorize with a masked array input:

In [180]: def foo(x):
     ...:     print(x)
     ...:     return 2*x
     ...:     

In [181]: np.vectorize(foo)(np.ma.masked_array([1,2,3],[True,False,True]))
1
1
2
3
Out[181]: 
masked_array(data=[--, 4, --],
             mask=[ True, False,  True],
       fill_value=999999)

In [182]: _.data
Out[182]: array([2, 4, 6])

Upvotes: 0

Related Questions