求助：基于statsmodels按code计算动态滚动窗口回归系数与截距失败

目标是使用statsmodels按code分组，基于动态滚动窗口（由每行的index_start和group_index定义窗口范围）计算回归系数与截距，现有固定窗口代码可运行，但动态窗口实现报错，且逐行处理效率低下。

样本数据 #

import pandas as pd
import statsmodels.api as sm
from statsmodels.regression.rolling import RollingOLS

d={'code':['A','A','A','A','B','B','B','C','C','C'],
   'Price':[10,8,11,13,1.1,1.3,1.5,100,101,104],
   'group_index':[1,2,3,4,1,2,3,1,2,3],
   'index_start':[1,1,2,2,1,1,2,1,1,1]}
df_dynamic_window=pd.DataFrame(data=d)

错误原因分析 #

• RollingOLS动态窗口失败：RollingOLS的window参数仅支持固定整数，无法接收逐行变化的序列值，直接传入x['group_index']-x['index_start']会触发类型错误。
• 逐行函数报错：groupby.apply默认将整个分组的DataFrame传入函数，而非逐行传入，导致函数中row变量实际是整个分组的DataFrame，引发索引匹配错误；同时直接引用全局DataFrame会大幅降低效率。

方法一：分组内高效处理（推荐） #

在每个code分组内部逐行计算，避免全局数据查询，提升执行效率：

def dynamic_window_regression(group):
    # 预添加截距项，减少重复计算
    group['const'] = 1
    results = []
    for idx, row in group.iterrows():
        # 筛选当前行对应的动态窗口
        window_data = group[(group['group_index'] >= row['index_start']) & 
                            (group['group_index'] <= row['group_index'])]
        # 校验窗口样本量（至少2个样本才能拟合带截距的OLS）
        if len(window_data) >= 2:
            y = window_data['Price']
            X = window_data[['const', 'group_index']]
            ols_result = sm.OLS(y, X).fit()
            results.append({
                'intercept': ols_result.params['const'],
                'coef_group_index': ols_result.params['group_index']
            })
        else:
            # 样本不足时填充NaN
            results.append({'intercept': None, 'coef_group_index': None})
    # 将结果合并回原分组
    group[['intercept', 'coef_group_index']] = pd.DataFrame(results, index=group.index)
    return group

# 应用到所有分组
result_df = df_dynamic_window.groupby('code', group_keys=True).apply(dynamic_window_regression)

优化点 #

• 仅在分组内部操作，避免全局数据查询
• 预添加截距项，减少重复计算开销
• 增加样本量校验，避免拟合报错

方法二：逐行处理（修正版） #

若需逐行调用函数，需使用df.apply(axis=1)而非groupby.apply，并提前构建分组字典避免重复查询：

# 预构建code分组字典，避免重复查询全局数据
code_groups = {code: group for code, group in df_dynamic_window.groupby('code')}

def regression_row(row):
    group = code_groups[row['code']]
    start = row['index_start']
    end = row['group_index']
    window_data = group[(group['group_index'] >= start) & (group['group_index'] <= end)]
    
    if len(window_data) >= 2:
        y = window_data['Price']
        X = sm.add_constant(window_data['group_index'])
        ols_result = sm.OLS(y, X).fit()
        return pd.Series([ols_result.params[0], ols_result.params[1]], 
                         index=['intercept', 'coef_group_index'])
    else:
        return pd.Series([None, None], index=['intercept', 'coef_group_index'])

# 逐行应用函数
df_dynamic_window[['intercept', 'coef_group_index']] = df_dynamic_window.apply(regression_row, axis=1)

大数据量性能建议 #

• 优先选择方法一，分组内循环比全局逐行查询效率更高
• 若窗口范围有固定规律（如index_start为固定偏移），可预计算窗口索引，避免每次循环筛选数据
• 超大规模数据可结合numba加速分组内循环逻辑

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