python Pandas : how to run multiple univariate regression by group

Question

假设我有一个 DataFrame，其中有一列 y 变量和许多列 x 变量。我希望能够运行 y 与 x1、y 与 x2 的多个单变量回归，.. ., 等等，并将预测存储回 DataFrame。我还需要通过组变量来执行此操作。

import statsmodels.api as sm
import pandas as pd

df = pd.DataFrame({
  'y': np.random.randn(20),
  'x1': np.random.randn(20), 
  'x2': np.random.randn(20),
  'grp': ['a', 'b'] * 10})

def ols_res(x, y):
    return sm.OLS(y, x).fit().predict()

df.groupby('grp').apply(ols_res) # This does not work

上面的代码显然是行不通的。我不清楚如何在 apply 遍历 x 列时正确地将固定的 y 传递给函数(x1 , x2, ...)。我怀疑可能有一个非常聪明的单线解决方案来做到这一点。有什么想法吗？

Answer 1

传递给 apply 的函数必须将 pandas.DataFrame 作为第一个参数。您可以将附加的关键字或位置参数传递给 apply，这些参数会传递给应用函数。因此，您的示例只需稍作修改即可使用。将 ols_res 更改为

def ols_res(df, xcols,  ycol):
    return sm.OLS(df[ycol], df[xcols]).fit().predict()

然后，你可以像这样使用groupby和apply

df.groupby('grp').apply(ols_res, xcols=['x1', 'x2'], ycol='y')

或者

df.groupby('grp').apply(ols_res, ['x1', 'x2'], 'y')

编辑

上面的代码不运行多个单变量回归。相反，它对每组运行一个多元回归。然而，通过(另一个)轻微的修改，它会。

def ols_res(df, xcols,  ycol):
    return pd.DataFrame({xcol : sm.OLS(df[ycol], df[xcol]).fit().predict() for xcol in xcols})

编辑 2

虽然上面的解决方案有效，但我认为下面的解决方案更像 pandas-y

import statsmodels.api as sm
import pandas as pd
import numpy as np

df = pd.DataFrame({
  'y': np.random.randn(20),
  'x1': np.random.randn(20), 
  'x2': np.random.randn(20),
  'grp': ['a', 'b'] * 10})

def ols_res(x, y):
    return pd.Series(sm.OLS(y, x).fit().predict())

df.groupby('grp').apply(lambda x : x[['x1', 'x2']].apply(ols_res, y=x['y']))

出于某种原因，如果我按原样定义 ols_res()，生成的 DataFrame 在索引中没有组标签。