Skip to content
Advertisement

Pandas rolling apply function to entire window dataframe

I want to apply a function to a rolling window. All the answers I saw here are focused on applying to a single row / column, but I would like to apply my function to the entire window. Here is a simplified example:

import pandas as pd
data = [ [1,2], [3,4], [3,4], [6,6], [9,1], [11,2] ]
df = pd.DataFrame(columns=list('AB'), data=data)

This is df:

    A   B
0   1   2
1   3   4
2   3   4
3   6   6
4   9   1
5   11  2

Take some function to apply to the entire window:

df.rolling(3).apply(lambda x: x.shape)

In this example, I would like to get something like:

    some_name   
0   NA  
1   NA  
2   (3,2)   
3   (3,2)   
4   (3,2)   
5   (3,2)   

Of course, the shape is used as an example showing f treats the entire window as the object of calculation, not just a row / column. I tried playing with the axis keyword for rolling, as well as with the raw keyword for apply but with no success. Other methods (agg, transform) do not seem to deliver either.

Sure, I can do this with a list comprehension. Just thought there is an easier / cleaner way of doing this.

Advertisement

Answer

Not with pd.DataFrame.rolling …. that function is applied iteratively to the columns, taking in a series of floats/NaN, and returning a series of floats/NaN, one-by-one. I think you’ll have better luck with your intuition….

def rolling_pipe(dataframe, window, fctn):
    return pd.Series([dataframe.iloc[i-window: i].pipe(fctn) 
                      if i >= window else None 
                      for i in range(1, len(dataframe)+1)],
                     index = dataframe.index) 

df.pipe(rolling_pipe, 3, lambda x: x.shape)
User contributions licensed under: CC BY-SA
6 People found this is helpful
Advertisement