python - Pandas - Using `.rolling()` on multiple columns

Question

Consider a pandas DataFrame which looks like the one below

      A     B     C
0  0.63  1.12  1.73
1  2.20 -2.16 -0.13
2  0.97 -0.68  1.09
3 -0.78 -1.22  0.96
4 -0.06 -0.02  2.18

I would like to use the function .rolling() to perform the following calculation for t = 0,1,2:

• Select the rows from t to t+2
• Take the 9 values contained in those 3 rows, from all the columns. Call this set S
• Compute the 75th percentile of S (or other summary statistics about S)

For instance, for t = 1 we have S = { 2.2 , -2.16, -0.13, 0.97, -0.68, 1.09, -0.78, -1.22, 0.96 } and the 75th percentile is 0.97.

I couldn't find a way to make it work with .rolling(), since it apparently takes each column separately. I'm now relying on a for loop, but it is really slow.

Do you have any suggestion for a more efficient approach?

Answer 1

One solution is to stack the data and then multiply your window size by the number of columns and slice the result by the number of columns. Also, since you want a forward looking window, reverse the order of the stacked DataFrame

wsize = 3
cols = len(df.columns)

df.stack(dropna=False)[::-1].rolling(window=wsize*cols).quantile(0.75)[cols-1::cols].reset_index(-1, drop=True).sort_index()

Output:

0    1.12
1    0.97
2    0.97
3     NaN
4     NaN
dtype: float64

In the case of many columns and a small window:

import pandas as pd
import numpy as np

wsize = 3
df2 = pd.concat([df.shift(-x) for x in range(wsize)], 1)
s_quant = df2.quantile(0.75, 1)

# Only necessary if you need to enforce sufficient data. 
s_quant[df2.isnull().any(1)] = np.NaN

Output: s_quant

0    1.12
1    0.97
2    0.97
3     NaN
4     NaN
Name: 0.75, dtype: float64