如何优化Pandas DataFrame与Scipy的大规模Pearson相关系数计算

一、核心优化：替换双重循环为向量化/批量运算 #

针对49262列的大规模数据集，Python双重循环的效率瓶颈无法接受，必须用底层优化的批量运算替代，以下是几种可行方案：

1. Numpy+Scipy 手动实现批量计算 #

通过矩阵运算直接推导相关系数和p值，底层依赖C实现的numpy运算，比Python循环快数个数量级，同时支持逐对删除缺失值：

import numpy as np
from scipy.stats import t
import pandas as pd

# 第一步：预处理过滤无效列（有效样本数<2的列直接排除）
valid_col_mask = df.notna().sum(axis=0) >= 2
valid_cols = df.columns[valid_col_mask]
df_valid = df[valid_cols]

# 计算每对列的共同非缺失样本数
mask = df_valid.notna().values
pairwise_n = mask.T @ mask  # 矩阵乘法批量计算所有列对的有效样本数

# 计算均值、标准差（仅用列内有效样本）
means = df_valid.mean(axis=0).values
stds = df_valid.std(axis=0, ddof=1).values

# 计算协方差矩阵与相关系数矩阵
centered_data = df_valid.sub(means, axis=1).values
cov_matrix = (centered_data.T @ centered_data) / (pairwise_n - 1)
corr_matrix = cov_matrix / np.outer(stds, stds)

# 计算p值（基于t分布，处理r=±1的边界情况避免除零）
r_squared = corr_matrix ** 2
df_degree = pairwise_n - 2
t_stat = corr_matrix * np.sqrt(df_degree / np.maximum(1 - r_squared, 1e-10))
p_matrix = 2 * (1 - t.cdf(np.abs(t_stat), df_degree))

# 转换为DataFrame格式
corr_df = pd.DataFrame(corr_matrix, index=valid_cols, columns=valid_cols)
p_df = pd.DataFrame(p_matrix, index=valid_cols, columns=valid_cols)

2. 用Pingouin库简化批量计算 #

Pingouin是统计分析专用库，内置pairwise_corr方法直接批量计算所有列对的相关系数和p值，自动处理缺失值，代码更简洁：

import pingouin as pg
import pandas as pd

# 预处理过滤无效列
valid_cols = df.columns[df.notna().sum(axis=0) >= 2]
df_valid = df[valid_cols]

# 批量计算所有列对的Pearson相关系数与p值（逐对删除缺失值）
corr_results = pg.pairwise_corr(df_valid, method='pearson', padjust=None)
# 结果是结构化DataFrame，包含'X'（列名1）、'Y'（列名2）、'r'（相关系数）、'p-val'（p值）等字段

3. 并行计算进一步提速 #

如果上述方法仍有压力，可通过并行框架拆分计算任务，利用多CPU核心：

Joblib并行处理列对

from joblib import Parallel, delayed
import numpy as np
from scipy.stats import pearsonr
import pandas as pd

# 预处理
valid_col_mask = df.notna().sum(axis=0) >= 2
valid_cols = df.columns[valid_col_mask]
df_valid_np = df[valid_cols].values
col_indices = list(range(len(valid_cols)))

# 定义单对列的计算函数
def compute_single_pair(i, j):
    x = df_valid_np[:, i]
    y = df_valid_np[:, j]
    valid_mask = ~np.isnan(x) & ~np.isnan(y)
    x_clean = x[valid_mask]
    y_clean = y[valid_mask]
    if len(x_clean) < 2:
        return (i, j, np.nan, np.nan)
    r, p = pearsonr(x_clean, y_clean)
    return (i, j, r, p)

# 仅生成上三角列对（避免重复计算对称项）
pairs = [(i, j) for i in col_indices for j in col_indices if i < j]

# 并行计算（n_jobs=-1使用全部CPU核心）
results = Parallel(n_jobs=-1, verbose=10)(delayed(compute_single_pair)(i,j) for i,j in pairs)

# 整理结果为对称矩阵
corr_matrix = np.full((len(valid_cols), len(valid_cols)), np.nan)
p_matrix = np.full((len(valid_cols), len(valid_cols)), np.nan)
for i, j, r, p in results:
    corr_matrix[i,j] = corr_matrix[j,i] = r
    p_matrix[i,j] = p_matrix[j,i] = p
np.fill_diagonal(corr_matrix, 1.0)
np.fill_diagonal(p_matrix, 0.0)

corr_df = pd.DataFrame(corr_matrix, index=valid_cols, columns=valid_cols)
p_df = pd.DataFrame(p_matrix, index=valid_cols, columns=valid_cols)

Dask处理超大规模内存数据

如果数据集大到无法装入内存，用Dask拆分数据块并行计算：

import dask.dataframe as dd
import numpy as np
from scipy.stats import t

# 转换为Dask DataFrame（根据CPU核心数设置分区数）
dask_df = dd.from_pandas(df, npartitions=8)

# 预处理过滤无效列
valid_col_mask = dask_df.notna().sum(axis=0).compute() >= 2
valid_cols = dask_df.columns[valid_col_mask]
dask_valid = dask_df[valid_cols]

# 计算相关系数矩阵
corr_matrix = dask_valid.corr(method='pearson').compute()

# 基于相关系数矩阵计算p值（同Numpy方案逻辑）
pairwise_n = dask_valid.notna().values.T @ dask_valid.notna().values
df_degree = pairwise_n - 2
r_squared = corr_matrix ** 2
t_stat = corr_matrix * np.sqrt(df_degree / np.maximum(1 - r_squared, 1e-10))
p_matrix = 2 * (1 - t.cdf(np.abs(t_stat), df_degree))

p_df = pd.DataFrame(p_matrix, index=valid_cols, columns=valid_cols)

二、整数索引vs列名：整数索引确实更快 #

直接用整数索引访问列（如df_valid_np[:, i]或df.iloc[:, i]）比列名索引（df[col_name]）快。原因是字符串列名需要哈希查找定位，而整数索引直接对应内存位置，在循环或批量索引场景下，速度差异会被放大。因此如果必须遍历列，优先用整数索引。

三、额外优化细节 #

• 预处理减载：先过滤有效样本数<2的列，直接减少计算量——比如若10%的列无效，计算量可减少约19%。
• 内存压缩：将数据类型转换为float32（若精度允许），减少内存占用，加快运算速度：df_valid = df_valid.astype('float32')。
• 避免重复计算：仅计算上三角矩阵的列对，利用相关系数的对称性，直接减少一半计算量。

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