如何在Polars DataFrame中按组计算线性回归残差（避免循环）

要避免循环拼接，直接用Polars的窗口函数（over）或者分组聚合后关联的方式就能实现，以下是两种可行方案：

方案一：窗口函数+批量映射（保留原行直接计算） #

这种方法直接在原DataFrame上新增残差列，无需额外关联操作：

1. 先定义线性回归拟合函数，输入分组后的X、Y数据，返回预测值：

import polars as pl
import numpy as np

def get_predicted_y(group: pl.DataFrame) -> pl.Series:
    # 构造带截距项的特征矩阵
    x = np.hstack([np.ones((len(group), 1)), group["X"].to_numpy().reshape(-1, 1)])
    y = group["Y"].to_numpy()
    # 最小二乘法求解系数
    coeffs = np.linalg.lstsq(x, y, rcond=None)[0]
    # 计算预测值
    y_pred = x @ coeffs
    return pl.Series(y_pred, name="y_pred")

2. 用map_batches结合over分组计算预测值，再生成残差列：

# 假设你的原始DataFrame名为df
df = df.with_columns(
    # 按GROUP1和GROUP2分组，对每组的X、Y计算预测值
    y_pred=pl.struct(["X", "Y"]).map_batches(get_predicted_y).over(["GROUP1", "GROUP2"])
).with_columns(
    # 计算残差：真实值 - 预测值
    residual=pl.col("Y") - pl.col("y_pred")
).drop("y_pred")  # 不需要预测值可以删除这列

方案二：分组聚合系数后关联计算 #

如果更关注性能，先分组计算回归系数再关联回原数据也是不错的选择：

1. 分组计算每组的截距和斜率：

coeff_df = df.group_by(["GROUP1", "GROUP2"]).agg(
    pl.struct(["X", "Y"]).map_batches(
        lambda g: np.linalg.lstsq(
            np.hstack([np.ones((len(g),1)), g["X"].to_numpy().reshape(-1,1)]),
            g["Y"].to_numpy(),
            rcond=None
        )[0]
    ).alias("coeffs")
).with_columns(
    intercept=pl.col("coeffs").list.get(0),
    slope=pl.col("coeffs").list.get(1)
).drop("coeffs")

2. 关联系数到原数据，计算残差：

df = df.join(coeff_df, on=["GROUP1", "GROUP2"]).with_columns(
    residual=pl.col("Y") - (pl.col("intercept") + pl.col("X") * pl.col("slope"))
).drop("intercept", "slope")

注意事项 #

• 上述代码默认拟合带截距的线性模型，如果不需要截距，去掉特征矩阵中的常数项即可；
• 依赖numpy的lstsq方法，无需额外安装机器学习库，轻量高效；
• 两种方法都会保留原DataFrame的所有行，仅新增residual列。

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