如何对Numpy数组的皮尔逊相关系数计算过程进行向量化优化

优化方案 #

你可以通过两步向量化操作完全消除Python层面的循环，运算速度可从几天缩短到几十分钟甚至更短：

1. 预计算所有样本行的z-score，避免原代码中内循环重复计算每行的标准化结果
2. 用批量矩阵乘法一次性计算所有两两向量的点积，直接得到完整的相关系数矩阵，底层调用高度优化的BLAS库，充分利用CPU的SIMD指令和多核心性能

优化后的代码如下：

import numpy as np

# L = 45
# M = 102272
# data 为形状 (M, L) 的float32数组

# 按行计算均值和标准差，keepdims保持维度方便广播
mean = np.mean(data, axis=1, keepdims=True)
std = np.std(data, axis=1, keepdims=True)
# 可选：避免标准差为0导致除0报错
std[std < 1e-8] = 1e-8
# 一次性得到所有行标准化后的z矩阵
z = (data - mean) / std

# 矩阵乘法直接计算所有两两相关系数
cmat = (z @ z.T) / L

注：原代码存在变量名笔误，初始化的矩阵名为cmat，循环内写为cmmat，优化后代码已修正该问题。

可选GPU加速 #

如果你的设备配备NVIDIA显卡，可安装cupy库，仅需极少量代码改动即可调用GPU完成运算，速度还能再提升数倍：

import cupy as cp

# 将数据转移到GPU显存
data_gpu = cp.array(data)
mean = cp.mean(data_gpu, axis=1, keepdims=True)
std = cp.std(data_gpu, axis=1, keepdims=True)
std[std < 1e-8] = 1e-8
z = (data_gpu - mean) / std
cmat_gpu = (z @ z.T) / L
# 结果转回CPU numpy数组
cmat = cmat_gpu.get()

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