pandas: computing rolling stats for every n-th step efficiently?

user_15673860 0
Viewed 44

I am writing a resampler which basically returns statistics of a dataframe (e.g. rolling mean) over various window sizes W. The base step size of my dataframe time series data is delta_t. If I just want to return the rolling statistic for every n-th step of the original dataframe, I currently sloppily do it like this:

df.rolling(W).mean().iloc[::n]

Of course this computes a lot of unnecessary fields and then just selects every n-th. So, is there an efficient way of just computing every n-th step rolling mean?

Thanks!

Best, JZ

rolling-computation pandas python
1 Answers

Let's call your original method f1:

def f1(df: pd.DataFrame, W: int, n: int) -> pd.DataFrame:
    return df.rolling(W).mean().iloc[::n]

It's pretty efficient and more importantly, very easy to understand. However, if your ultimate objective is performance, you can implement your own rolling algorithm:

def f2(df: pd.DataFrame, W: int, n: int) -> pd.DataFrame:
    # Convert to numpy array for speed
    arr = df.to_numpy()
    # Each row in `index` contains the indicies that
    # participate in the window. The windows are preselected
    # to those that we want in the final output.
    index = np.arange(0, len(arr), n)
    index = index[:, None] - range(W)
    
    # Calculating the mean
    result = arr[index].mean(axis=1)
    # Any window with negative indexes are assigned NaN
    result[(index < 0).any(axis=1)] = np.nan

    return pd.DataFrame(result, columns=df.columns, index=df.index[::n])

Validation:

df = pd.DataFrame(np.random.randint(1, 1000, size=(100_000, 5)), columns=list("abcde"))
W, n = 5, 10

assert f1(df1, W, n).equals(f2(df, W, n))

Benchmark:

%timeit f1(df, W, n)
19.9 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit f2(df, W, n)
5.41 ms ± 251 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Performance characteristics:

f1 runs in about constant time (~20ms on my computer), regardless of W and n. The performance of f2 varies based on the W / n ratio. The smaller this ratio, the faster it runs. However, if this ratio is sufficiently large, it may run slower than f1. The break-even point appears to be around 2.

Conclusion: performance is all about trade-offs. You get nothing for free. Go with f1 if it's fast enough for your use case as its simplicity and readability are hard to beat. Use f2 when W / n is small and performance is critical.

user_25389390
Related