Python pandas：如何按组进行多元线性回归

Question

假设我有一个叫做 DataFrame 的数据表，这个表里有一列是 y 变量，还有很多列是 x 变量。我想要对 y 和每一个 x 变量（比如 x1、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

上面的代码显然是不能工作的。我不太明白如何正确地把固定的 y 传递给这个函数，同时让 apply 在 x 列（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 的风格。

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 在索引中就没有组标签。