Pandas rolling apply function to entire window dataframe

Question

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.

Answer 1

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)

Answer 2

The argument supplied to your apply function is a Series with an index property containing start, stop and step properties.

RangeIndex(start=0, stop=2, step=1)

You can use this to query your data frame.

df = pd.DataFrame([('Sean', i) for i in range(1,11)], columns=['name', 'value'])

def func(series):
    view = df.iloc[series.index]
    # use view to do something...
    count = len(view[view.value.isin([1,2,8])])
    return count

df['count'] = df.value.rolling(2).apply(func)

There may be a more efficient way to do this but I'm not sure how.