pandas groupby shift не соблюдает группы

Question

У меня есть следующий DataFrame и произвольная функция

df = pd.DataFrame(
    {'grp': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3],
     'val': [0.80485036, 0.30698609, 0.33518013, 0.12214516, 0.66355629,
       0.71277808, 0.07193942, 0.97128731, 0.46351423, 0.81494857,
       0.82267912, 0.33043168, 0.55643, 0.63413976, 0.37998928, 0.54695376,
       0.99751999, 0.02726808, 0.2392102 , 0.93278521, 0.41905688]}
)

def myfunc(arr):
    return np.product(1+arr) - 1

Я вычисляю myfunc переход в группы:

df.groupby('grp')['val'].rolling(3).apply(myfunc)

grp    
1    0          NaN
     1          NaN
     2     2.149576
     3     0.958213
     4     1.492450
     5     2.197331
     6     2.054280
     7     2.619272
     8     2.092553
     9     4.236139
     10    3.841406
2    11         NaN
3    12         NaN
     13         NaN
     14    2.509898
     15    2.488528
     16    3.264265
     17    2.174331
     18    1.542845
     19    1.460438
     20    2.398822

Все хорошо. Теперь мне нужно сдвинуть скользящий cal c внутри группы на пять периодов.

df.groupby('grp')['val'].rolling(3).apply(myfunc).shift(-5)

grp    
1    0     2.197331
     1     2.054280
     2     2.619272
     3     2.092553
     4     4.236139
     5     3.841406
     6          NaN
     7          NaN
     8          NaN
     9     2.509898
     10    2.488528
2    11    3.264265
3    12    2.174331
     13    1.542845
     14    1.460438
     15    2.398822
     16         NaN
     17         NaN
     18         NaN
     19         NaN
     20         NaN
Name: val, dtype: float64

Что здесь происходит ?! Вся цель groupby - поддерживать границы между группами. Как (и почему ) pandas не соблюдает это. Должно быть:

grp    
1    0     2.197331
     1     2.054280
     2     2.619272
     3     2.092553
     4     4.236139
     5     3.841406
     6          NaN
     7          NaN
     8          NaN
     9          NaN
     10         NaN
2    11         NaN
3    12    2.174331
     13    1.542845
     14    1.460438
     15    2.398822
     16         NaN
     17         NaN
     18         NaN
     19         NaN
     20         NaN
Name: val, dtype: float64

Это похоже на серьезную ошибку в pandas. Я что-то упускаю? Как я могу заставить groupby делать groupby?

Answer 1

Проблема в том, что при разбиении на части код

df.groupby('grp')['val'].rolling(3).apply(myfunc).shift(-5)

эквивалентен

tmp = df.groupby('grp')['val'].rolling(3).apply(myfunc)
out = tmp.shift(-5)

Здесь tmp - это нормальный pd.Series. И, как вы теперь можете догадаться, out сдвигается в обычной серии без какой-либо группировки. И это ожидаемое поведение.

Чтобы получить желаемый результат, вы можете связать его с другой группой по:

(df.groupby('grp')['val'].rolling(3).apply(myfunc)
   .groupby('grp').shift(-5)        # extra groupby here 
)

, и все должно быть хорошо.