Python pandas:如何按组运行多个单变量回归

Question

假设我有DataFrame一列y变量和多列x变量.我希望能够对yvs x1,yvs 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

虽然,上面的解决方案有效,但我认为以下是更多的熊猫

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