Pandas，如何通过 groupby() / 条件函数调用应用更改数据框？

Question

我有一个超过 100 万行的巨大时间序列数据框。它有每月股票收益列表，我想创建一个新行来跟踪前 3 个月的滚动总和。数据框首先包含所有公司 A 行，然后是所有公司 B 行，然后是所有公司 C 行...

例如：

          date                COMNAM    PRC    RET
395 2017-02-28       GAS NATURAL INC 12.650  0.000
396 2017-03-31       GAS NATURAL INC 12.700  0.010
397 2017-04-28       GAS NATURAL INC 12.500 -0.016
398 2017-05-31       GAS NATURAL INC 12.700  0.016
399 2017-06-30       GAS NATURAL INC 12.925  0.024
400 2017-07-31       GAS NATURAL INC 12.950  0.002
401 2017-08-31       GAS NATURAL INC    nan    nan
402 1985-12-31                   NaN    nan    nan
403 1986-01-31  MOBILE NATIONAL CORP 11.625    nan
404 1986-02-28  MOBILE NATIONAL CORP 13.250  0.140
405 1986-03-31  MOBILE NATIONAL CORP 14.188  0.071
406 1986-04-30  MOBILE NATIONAL CORP 14.938  0.053
407 1986-05-30  MOBILE NATIONAL CORP 14.625 -0.021
408 1986-06-30  MOBILE NATIONAL CORP 12.688 -0.132
409 1986-07-31  MOBILE NATIONAL CORP 13.312  0.049
410 1986-08-29  MOBILE NATIONAL CORP 13.312  0.000
411 1986-09-30  MOBILE NATIONAL CORP 14.250  0.070
412 1986-10-31  MOBILE NATIONAL CORP 13.375 -0.061
413 1986-11-28  MOBILE NATIONAL CORP 13.375  0.000
414 1986-12-31  MOBILE NATIONAL CORP 12.375 -0.075

rolling() 函数将给我前 3 个月的总和，但这将包括每家公司第一个日期的前一个股票的最后回报。我觉得 groupby() 函数可能会有所帮助，但我有点坚持如何做。还是我想多了，有什么更好的方法我什至不需要 groupby？

Answer 1

计算前个月的滚动总和（没有 当前月份），对于两列感兴趣的列，来自当前组 行，定义如下函数：

def mySum(grp):
    return grp[['PRC', 'RET']].shift().rolling(3).sum() 

然后，要获得每个组（公司）的滚动总和，运行：

result = df.join(df[df.COMNAM.notnull()].groupby('COMNAM').apply(mySum)\
    .reset_index(level=0, drop=True).add_prefix('r'))

结果是当前 df 和调用结果之间的连接 每个组（公司）的上述功能。 中间结果的列名前加 r，以标记 滚动总和。

对于您的数据样本，结果是：

         date                COMNAM     PRC    RET    rPRC   rRET
0  2017-02-28       GAS NATURAL INC  12.650  0.000     NaN    NaN
1  2017-03-31       GAS NATURAL INC  12.700  0.010     NaN    NaN
2  2017-04-28       GAS NATURAL INC  12.500 -0.016     NaN    NaN
3  2017-05-31       GAS NATURAL INC  12.700  0.016  37.850 -0.006
4  2017-06-30       GAS NATURAL INC  12.925  0.024  37.900  0.010
5  2017-07-31       GAS NATURAL INC  12.950  0.002  38.125  0.024
6  2017-08-31       GAS NATURAL INC     NaN    NaN  38.575  0.042
7  1985-12-31                   NaN     NaN    NaN     NaN    NaN
8  1986-01-31  MOBILE NATIONAL CORP  11.625    NaN     NaN    NaN
9  1986-02-28  MOBILE NATIONAL CORP  13.250  0.140     NaN    NaN
10 1986-03-31  MOBILE NATIONAL CORP  14.188  0.071     NaN    NaN
11 1986-04-30  MOBILE NATIONAL CORP  14.938  0.053  39.063    NaN
12 1986-05-30  MOBILE NATIONAL CORP  14.625 -0.021  42.376  0.264
13 1986-06-30  MOBILE NATIONAL CORP  12.688 -0.132  43.751  0.103
14 1986-07-31  MOBILE NATIONAL CORP  13.312  0.049  42.251 -0.100
15 1986-08-29  MOBILE NATIONAL CORP  13.312  0.000  40.625 -0.104
16 1986-09-30  MOBILE NATIONAL CORP  14.250  0.070  39.312 -0.083
17 1986-10-31  MOBILE NATIONAL CORP  13.375 -0.061  40.874  0.119
18 1986-11-28  MOBILE NATIONAL CORP  13.375  0.000  40.937  0.009
19 1986-12-31  MOBILE NATIONAL CORP  12.375 -0.075  41.000  0.009

如果您想“忽略”NaN 值（将它们视为 0），请将 功能：

def mySum(grp):
    return grp[['PRC', 'RET']].fillna(0).shift().rolling(3).sum()