Pandas大DataFrame下，如何高效计算滚动1000行的相关矩阵？

我有一个大型DataFrame，所有列的值均为浮点数，索引为DatetimeIndex类型，示例数据如下：

col0 col1 col2 col3 col4 ...
2035-10-30  1.0  1.0  1.0  1.0  1.0 ...   
2035-10-31  1.0  1.0  1.0  1.0  1.0 ...
2035-11-01  1.0  1.0  1.0  1.0  1.0 ...
2035-11-02  1.0  1.0  1.0  1.0  1.0 ...
    ...     ...  ...  ...  ...  ... ...

需要针对每个索引（即每个日期），计算DataFrame最后1000行的相关矩阵，当前使用的循环代码如下：

for d in df.index:
    cor_mat = other_df[:d].tail(1000).corr()

其中df与other_df的维度和索引完全一致。由于DataFrame规模庞大，当前实现速度极慢。尝试过使用.rolling()方法但未成功，请问如何通过向量化操作或其他方式优化该过程？

直接循环计算每个窗口的相关矩阵效率极低，核心原因是重复计算了大量数据的统计量。可以利用滚动窗口的均值、协方差推导相关系数，完全用向量化操作实现，避免循环。

步骤1：计算滚动窗口的均值 #

先对other_df的每一列计算窗口大小为1000的滚动均值：

rolling_mean = other_df.rolling(window=1000, min_periods=1).mean()

min_periods=1用于处理窗口不足1000行的初始阶段（比如前999天），如果不需要保留这些初始值，可去掉该参数，此时前999行结果为NaN。

步骤2：计算滚动窗口的协方差矩阵 #

先计算每个元素减去对应窗口的均值，得到中心化数据：

centered = other_df - rolling_mean

再通过矩阵乘法结合滚动求和计算协方差矩阵：

import numpy as np

n_cols = other_df.shape[1]
# 初始化三维数组存储每个窗口的协方差矩阵，形状为(行数, 列数, 列数)
cov_matrices = np.zeros((len(other_df), n_cols, n_cols))

# 遍历列对计算滚动协方差
for i in range(n_cols):
    for j in range(i, n_cols):
        # 计算偏差乘积的滚动和，除以窗口大小-1得到无偏协方差
        cov = centered.iloc[:, i].mul(centered.iloc[:, j]).rolling(window=1000, min_periods=1).sum() / (other_df.rolling(window=1000, min_periods=1).count().iloc[:, 0] - 1)
        cov_matrices[:, i, j] = cov.values
        cov_matrices[:, j, i] = cov.values  # 协方差矩阵对称

步骤3：从协方差矩阵推导相关矩阵 #

相关系数等于协方差除以两个变量的标准差乘积，先计算每列的滚动标准差：

rolling_std = other_df.rolling(window=1000, min_periods=1).std()

再遍历每个窗口计算相关矩阵：

cor_matrices = np.zeros_like(cov_matrices)

for idx in range(len(other_df)):
    std = rolling_std.iloc[idx].values
    std_outer = np.outer(std, std)
    # 避免除以0，标准差为0时相关系数设为1（自身相关）或0（按需调整）
    mask = std_outer != 0
    cor_matrices[idx][mask] = cov_matrices[idx][mask] / std_outer[mask]
    np.fill_diagonal(cor_matrices[idx], 1.0)  # 对角线元素强制为1

步骤4：获取指定日期的相关矩阵 #

cor_matrices的第idx个元素，就是other_df截至other_df.index[idx]的最后1000行的相关矩阵，可直接按df的索引位置对应取值。

额外优化：Numba加速循环 #

如果列数较多，可用Numba编译协方差计算的循环，进一步提升速度：

from numba import njit

@njit
def compute_cov_matrices(centered_vals, window_size, min_periods):
    n_rows, n_cols = centered_vals.shape
    cov_matrices = np.zeros((n_rows, n_cols, n_cols))
    for i in range(n_cols):
        for j in range(i, n_cols):
            col_i = centered_vals[:, i]
            col_j = centered_vals[:, j]
            product = col_i * col_j
            for row in range(n_rows):
                start = max(0, row - window_size + 1)
                count = row - start + 1
                if count < min_periods:
                    cov = np.nan
                else:
                    sum_prod = product[start:row+1].sum()
                    cov = sum_prod / (count - 1)
                cov_matrices[row, i, j] = cov
                cov_matrices[row, j, i] = cov
    return cov_matrices

# 转换为numpy数组传入函数
centered_vals = centered.values
cov_matrices = compute_cov_matrices(centered_vals, 1000, 1)

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