Polars中Rolling OLS自定义函数正确性验证及y_hat实现方法

1. 自定义函数rolling_ols能否正确计算方程系数？ #

可以正确计算，参数对应关系完全匹配需求：

• 函数中sm.add_constant(np.column_stack((x, x**2)))生成的特征矩阵列顺序为常数项、x、x²
• RollingOLS拟合的模型为 y = c + b*x + a*x²，返回的params[0]对应常数项c，params[1]对应x的系数b，params[2]对应x²的系数a，与注释中y = ax² + bx + c的系数定义完全一致，函数逻辑正确。

2. 在Polars DataFrame中生成y_hat列的实现方法 #

结合Polars的向量化特性，我们先计算滚动回归的参数列，再基于参数列生成预测值y_hat（按需求计算x=12时的预测结果）：

import polars as pl
import numpy as np
import statsmodels.api as sm
from statsmodels.regression.rolling import RollingOLS

# y = ax² + bx + c
def rolling_ols(x, y, window_size):
    x_upper = sm.add_constant(np.column_stack((x, x**2)))
    model = RollingOLS(y, x_upper, window_size)
    results = model.fit()
    params = results.params
    # 将Statsmodels的参数结果转为Polars Series
    return (
        pl.Series(params.iloc[:, 0], name="c"),
        pl.Series(params.iloc[:, 1], name="b"),
        pl.Series(params.iloc[:, 2], name="a")
    )

# 窗口大小
size = 12

# 测试数据
test_df = pl.DataFrame(
    {
        'x': np.arange(100),
        'y': np.random.randn(100)
    }
)

# 获取滚动回归参数列
c_series, b_series, a_series = rolling_ols(test_df['x'].to_numpy(), test_df['y'].to_numpy(), window_size=size)

# 生成包含y_hat的最终DataFrame
df = test_df.with_columns(
    c_series,
    b_series,
    a_series,
    # 计算x=12时的预测值y_hat
    (pl.col("c") + pl.col("b") * 12 + pl.col("a") * (12 ** 2)).alias("y_hat")
).select('x', 'y', 'c', 'b', 'a', 'y_hat')

关键细节： #

• 前size-1行（即前11行）的参数值为NaN，这是RollingOLS的默认行为，因为滚动窗口未达到指定大小，无法完成回归计算
• 使用Polars的列表达式完成y_hat计算，保持向量化运算特性，符合Polars的性能优化设计

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