pandas rolling实现滚动成对线性回归报错的解决方案

假设有一个包含x、y两列的大型DataFrame，示例构造代码如下：

df = pd.DataFrame(
  {
    "x": np.linspace(0, 10, num=POINTS_NUM, endpoint=True), 
    "y": y
  }
)

需求是高效计算数据子集上滚动线性回归的斜率与截距，窗口大小以20为例。
原有实现采用rolling结合显式for循环的方案，代码如下：

regression_data = []
for window in df.rolling(window=int(20)):
    if window.shape[0] < 20:
        regression_data.append([None, None, None, None])
        continue

    lin_reg = linregress(window["x"], window["y"])
    regression_data.append(
        [window["x"].iloc[0], window["x"].iloc[-1], lin_reg.slope, lin_reg.intercept]
    )

该实现结果正确，但编码风格不够简洁，尝试直接调用Window对象的aggregate/apply方法复现逻辑时失败，测试代码如下：

def lin_reg_for_win(window):
    lin_reg = linregress(window["x"], window["y"])
    return [window.iloc[0], window.iloc[-1], lin_reg.slope, lin_reg.intercept]

df.rolling(2, method="table", min_periods=0).agg(lambda x: lin_reg_for_win(x))

运行后抛出错误：ValueError: Data must be 1-dimensional，需要符合pandas最佳实践、无显式for循环的高效实现方案。

rolling.apply/agg默认按逐列方式向自定义函数传入1维数组，无法同时获取窗口内x、y两列的完整数据，因此调用需要同时输入x、y序列的linregress时会触发维度错误。

线性回归的斜率、截距可以通过窗口内的统计量直接推导，无需逐窗口调用linregress，所有计算均可通过pandas底层C实现的滚动接口完成，性能比循环方案高1~2个数量级，完全适配大型DataFrame场景。
推导公式如下（窗口内样本量为n）：

• 斜率：slope = (n*sum(x*y) - sum(x)*sum(y)) / (n*sum(x²) - sum(x)²)
• 截距：intercept = (sum(y) - slope*sum(x)) / n

实现代码：

import numpy as np
import pandas as pd

WIN_SIZE = 20

# 预计算所有滚动统计量
roll_count = df["x"].rolling(WIN_SIZE).count()
roll_sum_x = df["x"].rolling(WIN_SIZE).sum()
roll_sum_y = df["y"].rolling(WIN_SIZE).sum()
roll_sum_xy = (df["x"] * df["y"]).rolling(WIN_SIZE).sum()
roll_sum_x2 = (df["x"] ** 2).rolling(WIN_SIZE).sum()

# 计算回归参数
slope = (roll_count * roll_sum_xy - roll_sum_x * roll_sum_y) / (roll_count * roll_sum_x2 - roll_sum_x ** 2)
intercept = (roll_sum_y - slope * roll_sum_x) / roll_count

# 组装结果，窗口长度不足时自动返回NaN，与原逻辑的None语义一致
result = pd.DataFrame({
    "x_start": df["x"].shift(WIN_SIZE - 1),
    "x_end": df["x"],
    "slope": slope,
    "intercept": intercept
})

该方案输出与原for循环方案完全一致，无任何Python层循环，性能最优。

如果需要沿用自定义函数的写法，需要指定method="table"让rolling传入整个窗口的二维数据，同时设置raw=True接收numpy数组减少类型转换开销，注意该方案本质还是逐窗口执行Python函数，性能和原循环接近，不推荐大数据场景使用：

from scipy.stats import linregress

WIN_SIZE = 20

def linreg_window(arr):
    x, y = arr[:, 0], arr[:, 1]
    if len(x) < WIN_SIZE:
        return [np.nan]*4
    reg = linregress(x, y)
    return [x[0], x[-1], reg.slope, reg.intercept]

result_apply = df.rolling(
    WIN_SIZE, 
    method="table", 
    min_periods=WIN_SIZE
).apply(linreg_window, raw=True, result_type="expand")
result_apply.columns = ["x_start", "x_end", "slope", "intercept"]

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