如何基于多个NumPy掩码数组逐元素应用函数生成新掩码数组

原函数及功能 #

我有一个接收4个单值输入、返回单个浮点输出的函数，用于估算(x,y)测量值在给定不确定性（独立高斯分布）下，落入x-y空间某限定区域的概率：

from scipy.stats import multivariate_normal
import numpy as np

grid_step = 0.25  # 单位为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

该函数通过预先生成的grid_pdf，将网格按输入的误差缩放、测量值偏移后，筛选出落在目标区域的网格点，加权求和后得到概率。

需求与当前实现 #

我需要将这个函数逐元素应用到4个形状相同但掩码不同的NumPy掩码数组，生成同形状的掩码数组。当前采用循环实现：

mask1 = np.array([[False, True, False],[True, True, True],[True, False, False]])
mask2 = np.array([[True, True, True],[True, True, False],[False, False, True]])
# 有效区域为[0,1]、[1,0]、[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)

# 合并掩码，仅处理所有输入都有效的位置
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, _ in np.ma.ndenumerate(all_mask):
    prob[i] = cal_prob(x[i], x_err[i], y[i], y_err[i])

但循环效率较低，希望找到无需循环的矢量化实现方法。

矢量化解决方案 #

利用NumPy的广播机制，我们可以一次性处理所有有效元素，避免循环：

# 1. 获取合并后的有效掩码：仅保留所有输入都未被掩码的位置
combined_mask = x.mask | x_err.mask | y.mask | y_err.mask
valid_indices = ~combined_mask

# 2. 提取所有有效元素，展平为一维数组
x_valid = x[valid_indices].ravel()
x_err_valid = x_err[valid_indices].ravel()
y_valid = y[valid_indices].ravel()
y_err_valid = y_err[valid_indices].ravel()

# 3. 扩展网格维度，实现与有效元素的广播运算
# grid_x形状为(M,N)，扩展为(1,M,N)，与(K,)的有效元素广播为(K,M,N)
grid_x_exp = grid_x[np.newaxis, :, :]
grid_y_exp = grid_y[np.newaxis, :, :]

# 计算每个有效元素对应的偏移缩放后的网格
x_grid = x_err_valid[:, np.newaxis, np.newaxis] * grid_x_exp + x_valid[:, np.newaxis, np.newaxis]
y_grid = y_err_valid[:, np.newaxis, np.newaxis] * grid_y_exp + y_valid[:, np.newaxis, np.newaxis]

# 4. 判断每个网格点是否在目标区域
PSB_grid = ((x_grid > 3) & (y_grid < 10) & (y_grid < 10**(0.23 * x_grid - 0.46)))

# 5. 计算每个有效元素的概率，对网格维度求和
PSB_prob_valid = np.sum(PSB_grid * grid_pdf[np.newaxis, :, :], axis=(1,2)) / norm_pdf

# 6. 创建结果掩码数组，填充有效结果
prob_vectorized = np.ma.zeros_like(x)
prob_vectorized[valid_indices] = PSB_prob_valid
prob_vectorized.mask = combined_mask

关键说明 #

• 广播机制：通过扩展grid_x/grid_y和有效元素的维度，让NumPy自动对每个有效元素生成对应的网格数据，无需手动循环。
• 掩码处理：先提取有效元素集中计算，再将结果回填到对应位置，保证掩码数组的结构与输入一致。
• 性能提升：矢量化运算利用NumPy的底层优化，相比循环能大幅提升处理大数组时的效率。

内容的提问来源于stack exchange，提问作者Ho-Hin Leung