Python pandas rolling_apply two column input into function

Question

Following on from this question Python custom function using rolling_apply for pandas, about using rolling_apply. Although I have progressed with my function, I am struggling to deal with a function that requires two or more columns as inputs:

Creating the same setup as before

import pandas as pd
import numpy as np
import random

tmp  = pd.DataFrame(np.random.randn(2000,2)/10000, 
                    index=pd.date_range('2001-01-01',periods=2000),
                    columns=['A','B'])

But changing the function slightly to take two columns.

def gm(df,p):
    df = pd.DataFrame(df)
    v =((((df['A']+df['B'])+1).cumprod())-1)*p
    return v.iloc[-1]

It produces the following error:

pd.rolling_apply(tmp,50,lambda x: gm(x,5))

  KeyError: u'no item named A'

I think it is because the input to the lambda function is an ndarray of length 50 and only of the first column, and doesn't take two columns as the input. Is there a way to get both columns as inputs and use it in a rolling_apply function.

Again any help would be greatly appreciated...

Answer 1

Not sure if still relevant here, with the new rolling classes on pandas, whenever we pass raw=False to apply, we are actually passing the series to the wraper, which means we have access to the index of each observation, and can use that to further handle multiple columns.

From the docs:

raw : bool, default None

False : passes each row or column as a Series to the function.

True or None : the passed function will receive ndarray objects instead. If you are just applying a NumPy reduction function this will achieve much better performance.

In this scenario, we can do the following:

### create a func for multiple columns
def cust_func(s):

    val_for_col2 = df.loc[s.index, col2] #.values
    val_for_col3 = df.loc[s.index, col3] #.values
    val_for_col4 = df.loc[s.index, col4] #.values
    
    ## apply over multiple column values
    return np.max(s) *np.min(val_for_col2)*np.max(val_for_col3)*np.mean(val_for_col4)
    

### Apply to the dataframe
df.rolling('10s')['col1'].apply(cust_func, raw=False)

Note that here we can still use all functionalities from pandas rolling class, which is particularly useful when dealing with time-related windows.

The fact that we are passing one column and using the entire dataframe feels like a hack, but it works in practice.

Answer 2

Looks like rolling_apply will try to convert input of user func into ndarray (http://pandas.pydata.org/pandas-docs/stable/generated/pandas.stats.moments.rolling_apply.html?highlight=rolling_apply#pandas.stats.moments.rolling_apply).

Workaround based on using aux column ii which is used to select window inside of manipulating function gm:

import pandas as pd
import numpy as np
import random

tmp = pd.DataFrame(np.random.randn(2000,2)/10000, columns=['A','B'])
tmp['date'] = pd.date_range('2001-01-01',periods=2000)
tmp['ii'] = range(len(tmp))            

def gm(ii, df, p):
    x_df = df.iloc[map(int, ii)]
    #print x_df
    v =((((x_df['A']+x_df['B'])+1).cumprod())-1)*p
    #print v
    return v.iloc[-1]

#print tmp.head()
res = pd.rolling_apply(tmp.ii, 50, lambda x: gm(x, tmp, 5))
print res