如何在pandas中对DataFrame分组后一次性添加回归模型计算列

你需求的功能可以通过groupby.apply组合实现，和R的tidyverse写法逻辑对齐，性能和可读性都最优，以下是对应实现代码：

和R参考代码输出完全对齐的实现（每个分组一行结果，同时保留模型对象和斜率系数） #

首先导入依赖：

import pandas as pd
import seaborn as sns
from scipy.stats import linregress

加载数据集：

# 加载鸢尾花数据集，和R的iris数据集完全一致
iris = sns.load_dataset('iris')

核心处理逻辑：

def group_regression(df):
    # 拟合y=petal_length, x=sepal_length的单变量线性回归
    reg_res = linregress(x=df["sepal_length"], y=df["petal_length"])
    # 返回自定义字段，可按需扩展返回回归的其他指标
    return pd.Series(
        {
            "lm_res": reg_res, # 存储完整回归模型对象，对应R的lm.res列
            "coe": reg_res.slope # 提取自变量系数，对应R的coe列
        }
    )

# 按物种分组执行回归，得到结果
iris_lm = iris.groupby("species", group_keys=False).apply(group_regression).reset_index()

输出iris_lm就能得到和R代码完全一致的结果，三个物种的coe值分别约为0.132、0.686、0.750。

常见场景扩展 #

场景1：不需要存储模型，只提取回归系数 #

可以简化函数逻辑，代码更轻量：

def get_reg_slope(df):
    return linregress(df["sepal_length"], df["petal_length"]).slope

iris_coe = iris.groupby("species").apply(get_reg_slope).reset_index(name="coe")

场景2：把回归系数回填到原始DataFrame的每一行 #

如果需要同一分组的所有行都挂载对应分组的回归系数，使用transform即可：

iris["reg_coe"] = iris.groupby("species")[["sepal_length", "petal_length"]].transform(get_reg_slope)

多变量回归的扩展 #

如果需要拟合多变量线性回归，替换拟合逻辑即可，整体框架不变，示例如下：

from sklearn.linear_model import LinearRegression

def multi_var_reg(df):
    # 多变量x：sepal_length、sepal_width，y：petal_length
    X = df[["sepal_length", "sepal_width"]]
    y = df["petal_length"]
    model = LinearRegression().fit(X, y)
    return pd.Series({
        "intercept": model.intercept_,
        "sepal_length_coe": model.coef_[0],
        "sepal_width_coe": model.coef_[1],
        "r2_score": model.score(X, y)
    })

iris_multi_reg = iris.groupby("species").apply(multi_var_reg).reset_index()

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