Is there a way for pandas rolling.apply to enter a dataframe into the function-CodePudding

I have code as below:

def fn(x):
    ...
df.rolling(length).apply(lambda x: fn(x))

I want the function to take a subset of the dataframe as input. Is this possible?

CodePudding user response：

One way can be the following:

df = pd.DataFrame({'a': range(5), 'b': range(5)})

def fn(x):
    return x.sum()

r = df.set_index('b').rolling(3)

[fn(x.reset_index()) for x in r]

Output:

[b    0
 a    0
 dtype: int64,
 b    1
 a    1
 dtype: int64,
 b    3
 a    3
 dtype: int64,
 b    6
 a    6
 dtype: int64,
 b    9
 a    9
 dtype: int64]

CodePudding user response：

Actually I found a way, while it doesn't really pass a portion of the df as an argument the df can be acquired from there.

def fn(x):
    df_frac = df.iloc[x.index.start: x.index.stop]
    return ...

df.rolling(length).apply(lambda x: fn(x))