pandas groupby：有效地构建多个列

Question

我有一个带有多索引的数据框，这样我就可以轻松地创建一个使用来自多列的数据作为输入的函数：

df = pd.DataFrame({('ALSN','VA.M'):range(5), ('ALSN','VB.M'):np.arange(5)+2,
                  ('ALVY','VA.M'):range(5), ('ALVY','VB.M'):np.arange(5)+20,
                  ('ALSN', 'VP.M'):np.arange(5)-10, ('ALVY','VP.M'):np.arange(5)-30,
                  ('ALGG', 'VP.M'):np.arange(5)/5.})

        ALGG    ALSN                    ALVY
        VP.M    VA.M    VB.M    VP.M    VA.M    VB.M    VP.M
0       0.0     0       2       -10     0       20      -30
1       0.2     1       3       -9      1       21      -29
2       0.4     2       4       -8      2       22      -28
3       0.6     3       5       -7      3       23      -27
4       0.8     4       6       -6      4       24      -26

我想对此进行过滤，然后将一个函数应用于 VA.M 和 VB.M 列：

df2 = g.filter(lambda z: z.name[-1] != 'G')
df2.groupby(level=0, axis=1).apply(lambda g: pd.Series(g[g.name]['VB.M']+g[g.name]['VA.M']))

    ALSN    ALVY
0   2       20
1   4       22
2   6       24
3   8       26
4   10      28

所以，到目前为止一切顺利。 但是，我真正想做的是构建一个函数，该函数将多列作为输入（如上），然后输出多列。因此，例如，它可以对 VA.M 和 VB.M 列求和，然后将平方根和立方根作为新列返回。

显然，我可以使用两个不同的应用函数（计算总和，然后计算平方根或立方根）来完成此操作，但我只想执行一次中间步骤（计算总和）。这可能吗？

此外，我希望将输出存储回数据框中，例如，列 ('ALSN', 'V2') 和 ('ALSN', 'V3') 可以在同时申请？还是我需要计算 V2 或 V3 列，然后获取结果数据帧并与原始数据帧合并？

Answer 1

这并不容易。

首先通过get_level_values 和boolean indexing 找到MultiIndex 的第一级值，然后通过slicers 选择：

lvl = df.columns.get_level_values(0).unique()
lvl = lvl[~lvl.str.contains('G')]
print (lvl)
Index(['ALSN', 'ALVY'], dtype='object')

idx = pd.IndexSlice
df2 = df.loc[:, idx[lvl, ['VA.M','VB.M']]]
print (df2)
  ALSN      ALVY     
  VA.M VB.M VA.M VB.M
0    0    2    0   20
1    1    3    1   21
2    2    4    2   22
3    3    5    3   23
4    4    6    4   24

对每个函数使用groupby，然后使用concat 和unstack：

print (pd.concat([df2.groupby(level=0, axis=1).apply(lambda x: (x**2).sum(axis=1)), 
                  df2.groupby(level=0, axis=1).apply(lambda x: (x**3).sum(axis=1))],
                  keys=('x^2','x^3')).unstack(0))

  ALSN      ALVY       
   x^2  x^3  x^2    x^3
0    4    8  400   8000
1   10   28  442   9262
2   20   72  488  10656
3   34  152  538  12194
4   52  280  592  13888

另一个非常相似的解决方案：

print (df2.groupby(level=0, axis=1).apply(lambda x: (x**2)))
  ALSN      ALVY     
  VA.M VB.M VA.M VB.M
0    0    4    0  400
1    1    9    1  441
2    4   16    4  484
3    9   25    9  529
4   16   36   16  576

print (df2.groupby(level=0, axis=1).apply(lambda x: (x**3)))
  ALSN      ALVY       
  VA.M VB.M VA.M   VB.M
0    0    8    0   8000
1    1   27    1   9261
2    8   64    8  10648
3   27  125   27  12167
4   64  216   64  13824

df21 = df2.groupby(level=0, axis=1).apply(lambda x: (x**2).sum(axis=1))
df22 = df2.groupby(level=0, axis=1).apply(lambda x: (x**3).sum(axis=1))
print (df21)
   ALSN  ALVY
0     4   400
1    10   442
2    20   488
3    34   538
4    52   592

print (df22)
   ALSN   ALVY
0     8   8000
1    28   9262
2    72  10656
3   152  12194
4   280  13888

print (pd.concat([df21,df22], keys=('x^2','x^3')).unstack(0))
  ALSN      ALVY       
   x^2  x^3  x^2    x^3
0    4    8  400   8000
1   10   28  442   9262
2   20   72  488  10656
3   34  152  538  12194
4   52  280  592  13888

通过评论编辑：

df = df2.groupby(level=0, axis=1).sum()
print (df)
   ALSN  ALVY
0     2    20
1     4    22
2     6    24
3     8    26
4    10    28

df3 = pd.concat([df ** 2, df ** 3], keys=('x^2','x^3'), axis=1)
df3.columns = df3.columns.swaplevel(0,1)
df3 = df3.sort_index(axis=1)
print (df3)
  ALSN       ALVY       
   x^2   x^3  x^2    x^3
0    4     8  400   8000
1   16    64  484  10648
2   36   216  576  13824
3   64   512  676  17576
4  100  1000  784  21952