如何使用 groupby + transform 而不是管道？

Question

假设我有一个这样的数据框

import pandas as pd
from scipy import stats

df = pd.DataFrame(
    {
        'group': list('abaab'),
        'val1': range(5),
        'val2': range(2, 7),
        'val3': range(4, 9)
    }
)

  group  val1  val2  val3
0     a     0     2     4
1     b     1     3     5
2     a     2     4     6
3     a     3     5     7
4     b     4     6     8

现在我想使用 vali 列中的两个（可能是所有对，所以我不想在任何地方硬编码列名）计算列 group 中每个组的线性回归。

基于 pipe

的潜在实现可能如下所示

def do_lin_reg_pipe(df, group_col, col1, col2):
    group_names = df[group_col].unique()
    df_subsets = []
    for s in group_names:
        df_subset = df.loc[df[group_col] == s]
        x = df_subset[col1].values
        y = df_subset[col2].values
        slope, intercept, r, p, se = stats.linregress(x, y)
        df_subset = df_subset.assign(
            slope=slope,
            intercept=intercept,
            r=r,
            p=p,
            se=se
        )
        df_subsets.append(df_subset)
    return pd.concat(df_subsets)

然后我可以使用

df_linreg_pipe = (
    df.pipe(do_lin_reg_pipe, group_col='group', col1='val1', col2='val3')
      .assign(p=lambda d: d['p'].round(3))
)

给出了想要的结果

  group  val1  val2  val3  slope  intercept    r    p   se
0     a     0     2     4    1.0        4.0  1.0  0.0  0.0
2     a     2     4     6    1.0        4.0  1.0  0.0  0.0
3     a     3     5     7    1.0        4.0  1.0  0.0  0.0
1     b     1     3     5    1.0        4.0  1.0  0.0  0.0
4     b     4     6     8    1.0        4.0  1.0  0.0  0.0

我不喜欢的是我必须循环遍历组，使用 append 然后 concat，所以我想我应该以某种方式使用 groupby和 transform 但我无法正常工作。函数调用应该类似于

df_linreg_transform = df.copy()
df_linreg_transform[['slope', 'intercept', 'r', 'p', 'se']] = (
    df.groupby('group').transform(do_lin_reg_transform, col1='val1', col2='val3')
)

问题是如何定义do_lin_reg_transform；我想要一些类似的东西

def do_lin_reg_transform(df, col1, col2):
    
    x = df[col1].values
    y = df[col2].values
    slope, intercept, r, p, se = stats.linregress(x, y)

    return (slope, intercept, r, p, se)

但是那 - 当然 - 崩溃 KeyError

KeyError: 'val1'

如何实现 do_lin_reg_transform 使其与 groupby 和 transform 一起工作？

Answer 1

因为您可以使用 groupby_transform 因为您需要额外的列来计算结果，所以我们的想法是使用 groupby_apply 和 map 将结果广播到每一行：

cols = ['slope', 'intercept', 'r', 'p', 'se']
lingress = lambda x: stats.linregress(x['val1'], x['val3'])

df[cols] = pd.DataFrame.from_records(df['group'].map(df.groupby('group').apply(lingress)))
print(df)

# Output
  group  val1  val2  val3  slope  intercept    r             p   se
0     a     0     2     4    1.0        4.0  1.0  9.003163e-11  0.0
1     b     1     3     5    1.0        4.0  1.0  0.000000e+00  0.0
2     a     2     4     6    1.0        4.0  1.0  9.003163e-11  0.0
3     a     3     5     7    1.0        4.0  1.0  9.003163e-11  0.0
4     b     4     6     8    1.0        4.0  1.0  0.000000e+00  0.0

Answer 2

转换旨在聚合单个列的结果。回归需要多个，所以你应该使用 apply.

如果需要，您可以将聚合定义为 return DataFrame 而不是 Series（因此结果不会减少）。为此，您需要确保索引是唯一的。然后 concat 返回结果，使其与索引对齐。如果有超过 1 个分组列，您将不会遇到任何问题。

def group_reg(gp, col1, col2):
    df = pd.DataFrame([stats.linregress(gp[col1], gp[col2])]*len(gp), 
                      columns=['slope', 'intercept', 'r', 'p', 'se'],
                      index=gp.index)
    return df

pd.concat([df, df.groupby('group').apply(group_reg, col1='val1', col2='val3')], axis=1)

  group  val1  val2  val3  slope  intercept    r             p   se
0     a     0     2     4    1.0        4.0  1.0  9.003163e-11  0.0
1     b     1     3     5    1.0        4.0  1.0  0.000000e+00  0.0
2     a     2     4     6    1.0        4.0  1.0  9.003163e-11  0.0
3     a     3     5     7    1.0        4.0  1.0  9.003163e-11  0.0
4     b     4     6     8    1.0        4.0  1.0  0.000000e+00  0.0

如何使用 groupby + transform 而不是管道？

How to use groupby + transform instead of pipe?

python

dataframe

pandas

pandas-groupby