如何高效在带Datetime索引的Pandas DataFrame上实现滑动窗口相关性计算

问题背景 #

已加载Yahoo Finance的VIX和GSPC收盘价数据，需要计算50天滑动窗口内两者的Pearson相关系数。原始的for循环实现逻辑可行但效率极低，需要更高效的实现方案。

原始数据结构示例：

VIX         GSPC
Date        
1990-01-02  17.240000   359.690002
1990-01-03  18.190001   358.760010
1990-01-04  19.219999   355.670013
1990-01-05  20.110001   352.200012
1990-01-08  20.260000   353.790009

原始低效实现代码：

data_size = len(x)
period = 50

df = pd.DataFrame()

for i in range(data_size-period):
    df.loc[i, "GSPC_VIX_corr"] = x[["GSPC", "VIX"]][i:i+period].corr().loc["GSPC", "VIX"]

高效实现方案 #

方案1：利用rolling().apply()自定义计算逻辑 #

通过pandas的滚动窗口对象结合自定义函数，实现向量化计算，避免逐次切片的开销：

period = 50

def rolling_corr(window):
    # window是形状为(period, 2)的二维数组，对应窗口内的VIX和GSPC数据
    return window[:,0].corr(window[:,1])

# 对数据滚动窗口应用自定义函数，设置min_periods=period确保只有完整窗口才计算
rolling_corr_series = x.rolling(window=period, min_periods=period).apply(rolling_corr, raw=True)
# 提取结果（两列计算的相关系数对称，取任意一列即可）
result = rolling_corr_series["GSPC"].dropna()

注意：设置raw=True让窗口以numpy数组形式传入函数，比默认的DataFrame形式更快。

方案2：基于相关系数公式的向量化计算（性能最优） #

Pearson相关系数的计算公式为：
$$\text{corr}(X,Y) = \frac{\text{cov}(X,Y)}{\text{std}(X) \times \text{std}(Y)}$$

利用pandas内置的滚动协方差、滚动标准差函数，直接通过向量化操作计算，性能比apply()更优：

period = 50

# 计算滚动协方差
rolling_cov = x["GSPC"].rolling(window=period).cov(x["VIX"])
# 计算GSPC的滚动标准差
rolling_std_gspc = x["GSPC"].rolling(window=period).std()
# 计算VIX的滚动标准差
rolling_std_vix = x["VIX"].rolling(window=period).std()
# 计算滚动相关系数
rolling_corr = rolling_cov / (rolling_std_gspc * rolling_std_vix)
# 去除前period-1个空值
result = rolling_corr.dropna()

方案对比 #

• 原始for循环：时间复杂度O(N*K)（N为数据量，K为窗口大小），性能最差，数据量大时耗时极长。
• 方案1：rolling().apply()：时间复杂度O(N)，但自定义函数仍有一定开销，比循环快数倍。
• 方案2：公式化向量化计算：完全利用pandas的C级优化操作，性能最优，比方案1还要快2-3倍。

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