Best way to compute a rolling median absolution deviation (MAD) in polars (python)

Question

What would be the most efficient way to compute, in polars (python-polars), a rolling median absolute deviation (MAD) over a given period over a given "groupby key" ?

The context/the dataframe

I have the following dataframe:

• it contains 3 columns
  • 1 datetime column (type Datetime)
  • 1 group column (type UInt32)
  • 1 value column (type UInt32)
• the dataframe contains ~13 millions rows
  • it contains some values from 1500 distinct groups at the granularity of 5mn over 1 month
  • 5mn data points for a full month is around 8640 points
  • 8640 points for 1500 groups => 8640 * 1500 ~= 13 millions data points

┌─────────────────────┬───────┬───────┐
│ datetime            ┆ group ┆ value │
│ ---                 ┆ ---   ┆ ---   │
│ datetime[μs]        ┆ u32   ┆ u32   │
╞═════════════════════╪═══════╪═══════╡
│ 2023-04-17 14:50:00 ┆ 1     ┆ 8     │
│ 2023-04-17 14:50:00 ┆ 2     ┆ 4     │
│ 2023-04-17 14:50:00 ┆ 3     ┆ 14893 │
│ 2023-04-17 14:50:00 ┆ 4     ┆ 15    │
│ …                   ┆ …     ┆ …     │
│ 2023-05-17 14:45:00 ┆ 1497  ┆ 3     │
│ 2023-05-17 14:45:00 ┆ 1498  ┆ 8193  │
│ 2023-05-17 14:45:00 ┆ 1499  ┆ 42    │
│ 2023-05-17 14:45:00 ┆ 1500  ┆ 15    │
└─────────────────────┴───────┴───────┘

The processing

I now want to compute a rolling median absolute deviation over a week of data (i.e. 2016 observations as we are working at 5mn granularity) over every group.

I am currently able to compute very efficiently a rolling standard deviation with such kind of "expression":

(
    df.
    with_columns(
        pl.col("value").rolling_std(2016).over("group").alias("std")
    )
)

It gives me an answer in less than a second!

Unfortunately, the standard deviation is too sensitive to "outliers" and we would need to switch to a median absolute deviation (i.e. take the median of the absolute distance to the median instead of the mean of the distance to the mean).

The current approach

I gave a try with a rolling_apply + the computation of the median from the returned series but it is too slow (it is still processing after more than 5 minutes...compared to the 1s of the rolling_std).

(
    df
    .with_columns(
        pl.col("value").rolling_apply(lambda s: (s-s.median()).median(), window_size=2016).over("group").alias("mad")
    )
)

Improving the performance

Is there a way to do it more efficiently and to reach the same performance of what we can have with rolling_std ?

(i understand that computing the mad can be slower than computing a standard deviation...but i mean having the same kind of order of magnitude...i.e. to be able to compute in several seconds).

UPDATE

As noticed by @Dean Mac Gregor, the original lambda i've put in The current approach does not give the expected result (it returns always 0).
The proper lambda to have the expected result is rather lambda s: (s-s.median()).abs().median() (the call to abs() was missing).
The performance issue remains.
How to compute it more efficiently ?

Answer 1

I think you can use .cumulative_eval() for this.

s = pl.element().tail(2016)

df.with_columns(mad = 
   pl.col("value")
     .cumulative_eval((s - s.median()).abs().median(), parallel=True)
     .over("group")
)

You will get all the values, as if you supplied min_periods=1 to the rolling method.

Answer 2

Let's start with some sample data

import polars as pl
import numpy as np
from datetime import datetime
np.random.seed(0)
df=pl.DataFrame({'datetime':pl.date_range(datetime(2023,4,17,14,50), 
                                          datetime(2023,5,17,14,45), 
                                          '5m',eager=True)}) \
    .with_columns(group=pl.lit([[x for x in range(5)]])).explode('group') \
    .with_columns(value=pl.lit(np.random.randint(0,10000,size=43200)))

then your mad formula is

df=(
    df
    .with_columns(
        pl.col("value").rolling_apply(lambda s: (s-s.median()).abs().median(), window_size=2016).over("group").alias("mad")
    )
)

This takes about 5.8sec for me.

You can use groupby_rolling with an agg to get the functionality of rolling_* functions.

That'll look like this:

(
    df
        .groupby_rolling('datetime', period='7d',by='group')
        .agg(
             gbrmad=(pl.col('value')-pl.col('value').median()).abs().median()
             )
)

this takes 1 sec.

I notice that your lambda function returns an int and more specifically that it's always lower than what mine returns. If that is intended and not just an after effect then you can do:

(
    df
        .groupby_rolling('datetime', period='7d',by='group')
        .agg(
             gbrmad=(pl.col('value')-pl.col('value').median()).abs().median().floor().cast(pl.Int64())
             )
)

Lastly, you're doing a fixed window size of 2016 instead of explicitly a week. As such, you get nulls in the early part of the dataset. If you want to keep that behavior then we can do a when like this.

(
    df
    .groupby_rolling('datetime', period='7d',by='group')
    .agg(
            gbrmad=pl.when(pl.col('value').count()==2016)
                    .then((pl.col('value')-pl.col('value').median()).abs().median().floor().cast(pl.Int64()))
            ))

To check results:

df.join(
    df
    .groupby_rolling('datetime', period='7d',by='group')
    .agg(
            gbrmad=pl.when(pl.col('value').count()==2016)
                    .then((pl.col('value')-pl.col('value').median()).abs().median().floor().cast(pl.Int64()))
            ),
    on=['datetime','group']).filter(pl.col('mad')!=pl.col('gbrmad'))

0 rows where the lambda based mad calc is different with about a 5x improvement in speed.