熊猫，如何使用groupby() /条件函数调用来应用更改数据？

Question

我有一个有超过100万行的大型timeseries dataframe。它有每月股票回报的列表，我想创建一个新的行来跟踪前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()-function将给出前3个月的总和，但这将包括每个公司最初日期的前一批股票的最后回报。我觉得groupby()函数可能会帮上忙，但我还是想不出该怎么做。还是我想得太多了，还有什么更好的方法让我根本不需要群人呢？

Answer 1

要计算前3个月的滚动和(当前月份没有的)，需要从当前的行组中计算感兴趣的两列，定义以下函数：

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值(将其视为)，请将函数更改为：

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