如何优雅地对同形NumPy矩阵应用统计函数？

这题我太熟悉了！想要高效又优雅地处理这种对应位置逐组统计映射的需求，核心思路是利用NumPy的向量化操作，避免写繁琐的嵌套循环，同时精准控制输出形状。下面给你两种实用方案，按需选择：

方案一：用np.vectorize快速封装（简洁易读） #

如果你不想手动推导公式，直接复用scipy.stats.pearsonr的话，可以用numpy.vectorize把单组计算的函数“广播”到整个多维数组上：

import numpy as np
from scipy.stats import pearsonr

# 定义提取皮尔逊相关系数的辅助函数
def extract_corr(x, y):
    # pearsonr返回(相关系数, p值)，我们取第一个元素
    return pearsonr(x, y)[0]

# 向量化封装，指定输入输出的形状签名：两个(n,)数组输出一个标量
vectorized_corr = np.vectorize(extract_corr, signature='(n),(n)->()')

# 直接传入两个形状为(16,3,20)的数组，输出自动是(16,3)
correlations = vectorized_corr(data["histograms"][0], regions_hist[0])

这个方案的优点是代码直观，几乎不需要额外思考，完全贴合你想要的FancyMapping式调用体验。

方案二：手动向量化计算（性能拉满） #

如果你的数据量很大，追求极致性能，推荐直接用NumPy的广播和数组运算手动实现皮尔逊相关系数公式。这种方法完全避免了Python层面的循环，效率比vectorize高很多：

import numpy as np

x = data["histograms"][0]
y = regions_hist[0]

# 计算每组的均值（保持维度方便广播）
x_mean = x.mean(axis=2, keepdims=True)
y_mean = y.mean(axis=2, keepdims=True)

# 计算协方差（自由度为1，和pearsonr保持一致）
cov = ((x - x_mean) * (y - y_mean)).sum(axis=2) / (x.shape[2] - 1)

# 计算每组的标准差（自由度为1）
x_std = x.std(axis=2, ddof=1)
y_std = y.std(axis=2, ddof=1)

# 计算皮尔逊相关系数，自动得到(16,3)的结果
correlations = cov / (x_std * y_std)

验证形状与后续使用 #

两种方案得到的correlations形状都是(16,3)，完全符合你的需求。之后你可以直接做布尔判断：

mask = correlations > 0.42524395175128987

得到的mask也是(16,3)的布尔数组，方便后续的筛选或索引操作。

内容的提问来源于stack exchange，提问作者zar3bski