Python中逐网格点相关系数计算的代码优化问询

我需要在Python中计算每个网格点的皮尔逊相关系数，两个变量variable1[t,y,x]和variable2[t,y,x]维度一致，最终要得到[y,x]维度的结果。因为变量存在不规则缺失值，我写了一段嵌套循环的代码，但大量的判断和循环导致运行速度太慢，想找优化方法。

原代码如下：

def corr_grid3D(array1,array2):
    dims = array1.shape
    corr = np.empty([dims[1],dims[2]]); p = np.empty([dims[1],dims[2]])

    for i in range(dims[1]):
        for j in range(dims[2]):
            logic1 = ~np.isnan(array1[:,i,j]); logic2 = ~np.isnan(array2[:,i,j])
            if sum(logic1) == sum(logic2) and sum(logic1) > 5 and sum(logic2) > 5:
                corr[i,j], p[i,j] = stats.pearsonr(array1[logic1,i,j],array2[logic2,i,j])
            else:
                corr[i,j], p[i,j] = np.nan, np.nan
    return corr, p

优化思路与实现 #

1. 核心问题：嵌套循环拖慢速度 #

原代码的嵌套循环是性能瓶颈，尤其是当网格的y、x维度很大时，循环次数会达到y*x次，每次循环还要做缺失值判断、求和、调用pearsonr，开销极大。优化的核心是用numpy向量化操作替代循环，把批量计算一次性完成。

2. 统一处理缺失值掩码 #

首先创建一个共同的掩码，标记两个变量在同一时刻都非缺失的位置，这样就能一次性得到所有网格点的有效样本数：

mask = ~np.isnan(array1) & ~np.isnan(array2)
valid_counts = mask.sum(axis=0)  # 每个[y,x]网格点的有效样本数

这一步替代了原循环里的logic1、logic2以及求和判断，速度提升明显。

3. 向量化计算皮尔逊相关系数 #

根据皮尔逊相关系数的公式，我们可以拆解为向量化步骤计算，避免循环调用stats.pearsonr：

• 先把缺失值替换为0（后续计算会被掩码过滤，不影响结果）
• 计算每个网格点的有效数据均值
• 计算离均差、协方差、标准差，最终得到相关系数

4. 批量计算p值 #

利用t分布的性质，批量计算每个有效网格点的p值，避免循环调用统计函数。

完整优化代码 #

import numpy as np
from scipy.stats import distributions

def corr_grid3D_optimized(array1, array2):
    # 校验输入维度一致性
    assert array1.shape == array2.shape, "两个输入数组维度必须一致"
    
    # 生成共同掩码：仅保留两个变量都非缺失的时刻
    mask = ~np.isnan(array1) & ~np.isnan(array2)
    # 统计每个网格点的有效样本数
    valid_counts = mask.sum(axis=0)
    
    # 替换缺失值为0，方便后续向量化运算
    array1_masked = np.where(mask, array1, 0)
    array2_masked = np.where(mask, array2, 0)
    
    # 计算每个网格点的有效数据均值
    mean1 = array1_masked.sum(axis=0) / valid_counts
    mean2 = array2_masked.sum(axis=0) / valid_counts
    
    # 计算离均差（仅有效数据参与计算）
    dev1 = array1_masked - mean1[np.newaxis, ...]
    dev2 = array2_masked - mean2[np.newaxis, ...]
    
    # 计算协方差和标准差
    cov = (dev1 * dev2 * mask).sum(axis=0) / (valid_counts - 1)
    std1 = np.sqrt((dev1**2 * mask).sum(axis=0) / (valid_counts - 1))
    std2 = np.sqrt((dev2**2 * mask).sum(axis=0) / (valid_counts - 1))
    
    # 计算相关系数
    corr = cov / (std1 * std2)
    # 对有效样本数不足5的网格点设为NaN
    corr[valid_counts <= 5] = np.nan
    
    # 批量计算p值
    p = np.full_like(corr, np.nan)
    valid_idx = valid_counts > 5
    n = valid_counts[valid_idx]
    r_vals = corr[valid_idx]
    
    # 避免r=±1时的除以0问题，做一个小范围截断
    r_vals = np.clip(r_vals, -0.999999, 0.999999)
    # 计算t统计量
    t = r_vals * np.sqrt((n - 2) / (1 - r_vals**2))
    # 双尾p值
    p[valid_idx] = 2 * distributions.t.sf(np.abs(t), n - 2)
    
    return corr, p

优化效果 #

• 完全消除嵌套循环，利用numpy的向量化特性，计算速度会提升几个数量级，网格规模越大，提升越明显。
• 减少了冗余的缺失值判断和求和操作，代码逻辑更简洁。
• 批量计算相关系数和p值，避免了循环调用统计函数的额外开销。

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