как генерировать динамические столбцы в пандах

Question

У меня есть следующий фрейм данных в пандах

code    tank     nozzle_1   nozzle_2    nozzle_var
123     1        1          1           10   
123     1        2          2           12
123     2        1          1           10   
123     2        2          2           12

Я хочу рассчитать кумулятивную сумму группировки столбцов nozzle_1 и nozzle_2 по резервуару.Вот мой желаемый фрейм данных.

code    tank     nozzle_1   nozzle_2    nozzle_var   nozzle_1_cumsum   nozzle_2_cumsum
123     1        1          1           10           1                 1 
123     1        2          2           12           3                 3
123     2        1          1           10           1                 1
123     2        2          2           12           3                 3

Я получаю nozzle_1 и nozzle_2 из следующего кода в pandas

cols= df.columns[df.columns.str.contains(pat='nozzle_\d+$', regex=True)] 

Как рассчитать сумму по приведенному выше списку столбцов

Answer 1

Как насчет этого причудливого решения:

cols= df.columns[df.columns.str.contains(pat='nozzle_\d+$', regex=True)] 
df.assign(**df.groupby('tank')[cols].agg(['cumsum'])\
              .pipe(lambda x: x.set_axis(x.columns.map('_'.join), axis=1, inplace=False)))

Вывод:

   tank  nozzle_1  nozzle_2  nozzle_var  nozzle_1_cumsum  nozzle_2_cumsum
0     1         1         1          10                1                1
1     1         2         2          12                3                3
2     2         1         1          10                1                1
3     2         2         2          12                3                3

---

По шагам:

df_cumsum = df.groupby('tank')[cols].agg(['cumsum'])
df_cumsum.columns = df_cumsum.columns.map('_'.join)
pd.concat([df, df_cumsum], axis=1)

Вывод:

   tank  nozzle_1  nozzle_2  nozzle_var  nozzle_1_cumsum  nozzle_2_cumsum
0     1         1         1          10                1                1
1     1         2         2          12                3                3
2     2         1         1          10                1                1
3     2         2         2          12                3                3

Answer 2

import pandas as pd

d = {
    "tank":[1,1,1,2],
    "nozzle_1":[1,2,1,2],
    "nozzle_2":[1,2,1,2]
}

df = pd.DataFrame(d)
columns = df.columns[df.columns.str.match("nozzle_\d+$")]

df[columns + "_cumsum"] = \
    df.groupby("tank")[columns].cumsum(axis=0)
df

Выход

    tank    nozzle_1    nozzle_2    nozzle_1_cumsum nozzle_2_cumsum
0   1   1   1   1   1
1   1   2   2   3   3
2   1   1   1   4   4
3   2   2   2   2   2