Polars - How to compute rolling ewm grouped by column?

user_3781180 0
Viewed 75

What's the right way to perform a groupby + rolling aggregate operation in polars? For some reason performing an ewm_mean over a rolling groupby gives me the list of all the ewm's rolling by time. For example take the dataframe below:

shape: (10, 3)
┌─────────────────────┬────────┬───────────┐
│ ts                  ┆ symbol ┆ signal_0  │
│ ---                 ┆ ---    ┆ ---       │
│ datetime[ns]        ┆ cat    ┆ f64       │
╞═════════════════════╪════════╪═══════════╡
│ 2022-02-14 09:20:00 ┆ A      ┆ -1.704301 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-02-14 09:20:00 ┆ AA     ┆ -1.181743 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-02-14 09:50:00 ┆ A      ┆ 1.040125  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-02-14 09:50:00 ┆ AA     ┆ 0.776798  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ ...                 ┆ ...    ┆ ...       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-02-14 11:20:00 ┆ A      ┆ 2.088835  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-02-14 11:20:00 ┆ AA     ┆ 1.741544  │
└─────────────────────┴────────┴───────────┘

Here, I want to group by symbol and get the rolling mean for signal_0 at every timestamp. Unfortunately this doesn't work:

portfolios.groupby_rolling("ts", by="symbol", period="1d").agg(
    pl.col("signal_0").ewm_mean(half_life=0.1).alias(f"signal_0_mean")
)

^The above yields:

shape: (10, 3)
┌────────┬─────────────────────┬─────────────────────────────────────┐
│ symbol ┆ ts                  ┆ signal_0_mean                       │
│ ---    ┆ ---                 ┆ ---                                 │
│ cat    ┆ datetime[ns]        ┆ list[f64]                           │
╞════════╪═════════════════════╪═════════════════════════════════════╡
│ A      ┆ 2022-02-14 09:20:00 ┆ [-1.704301]                         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 09:50:00 ┆ [-1.704301, 1.037447]               │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 10:20:00 ┆ [-1.704301, 1.037447, 1.933809]     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 10:50:00 ┆ [-1.704301, 1.037447, ... 2.0732... │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ...    ┆ ...                 ┆ ...                                 │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 10:50:00 ┆ [-1.181743, 0.774887, ... 1.6235... │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 11:20:00 ┆ [-1.181743, 0.774887, ... 1.7414... │
└────────┴─────────────────────┴─────────────────────────────────────┘

Here is a mock polars dataframe to reproduce the issue:

pl.DataFrame([{'ts': datetime.datetime(2022, 2, 14, 9, 20),
  'symbol': 'A',
  'signal_0': -1.7043013917942946},
 {'ts': datetime.datetime(2022, 2, 14, 9, 20),
  'symbol': 'AA',
  'signal_0': -1.1817425250668494},
 {'ts': datetime.datetime(2022, 2, 14, 9, 50),
  'symbol': 'A',
  'signal_0': 1.040124717732041},
 {'ts': datetime.datetime(2022, 2, 14, 9, 50),
  'symbol': 'AA',
  'signal_0': 0.776797510847028},
 {'ts': datetime.datetime(2022, 2, 14, 10, 20),
  'symbol': 'A',
  'signal_0': 1.934685674908703},
 {'ts': datetime.datetime(2022, 2, 14, 10, 20),
  'symbol': 'AA',
  'signal_0': 1.480892267491577},
 {'ts': datetime.datetime(2022, 2, 14, 10, 50),
  'symbol': 'A',
  'signal_0': 2.0734176892015395},
 {'ts': datetime.datetime(2022, 2, 14, 10, 50),
  'symbol': 'AA',
  'signal_0': 1.6236975034431897},
 {'ts': datetime.datetime(2022, 2, 14, 11, 20),
  'symbol': 'A',
  'signal_0': 2.0888349547293954},
 {'ts': datetime.datetime(2022, 2, 14, 11, 20),
  'symbol': 'AA',
  'signal_0': 1.741544172734057}])

If I wanted to do this in pandas, I would write:

df.set_index(["ts", "symbol"]).groupby(level=1)["signal_0"].transform(
    lambda x: x.ewm(halflife=10).mean()
)

Which would yield:

ts                   symbol
2022-02-14 09:20:00  A        -1.704301
                     AA       -1.181743
2022-02-14 09:50:00  A        -0.284550
                     AA       -0.168547
2022-02-14 10:20:00  A         0.507021
                     AA        0.419785
2022-02-14 10:50:00  A         0.940226
                     AA        0.752741
2022-02-14 11:20:00  A         1.202843
                     AA        0.978820
Name: signal_0, dtype: float64
python-polars pandas python
2 Answers

You were close. Since ewm_mean produces an estimate for each observation in each window, you simply need to specify that you want the last calculated value in each rolling window.

(
    portfolios
    .groupby_rolling("ts", by="symbol", period="1d")
    .agg(
        pl.col("signal_0").ewm_mean(half_life=10).last().alias(f"signal_0_mean")
    )
    .sort(['ts', 'symbol'])
)

shape: (10, 3)
┌────────┬─────────────────────┬───────────────┐
│ symbol ┆ ts                  ┆ signal_0_mean │
│ ---    ┆ ---                 ┆ ---           │
│ str    ┆ datetime[μs]        ┆ f64           │
╞════════╪═════════════════════╪═══════════════╡
│ A      ┆ 2022-02-14 09:20:00 ┆ -1.704301     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 09:20:00 ┆ -1.181743     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 09:50:00 ┆ -0.2845       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 09:50:00 ┆ -0.168547     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 10:20:00 ┆ 0.507021      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 10:20:00 ┆ 0.419785      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 10:50:00 ┆ 0.940226      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 10:50:00 ┆ 0.752741      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 11:20:00 ┆ 1.202843      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 11:20:00 ┆ 0.9788        │
└────────┴─────────────────────┴───────────────┘
user_185598750

This is a solution, but I'm not sure if it's the most performant one:

portfolios.groupby("symbol")
.agg([pl.col("ts"), pl.col("signal_0")])
.with_column(pl.col("signal_0").arr.eval(pl.element().ewm_mean(half_life=10)).alias("signal_0_mean"))
.explode(["ts", "signal_0", "signal_0_mean"])
.sort(["ts", "symbol"])

Explanation:

  1. We first group by symbol to aggregate all the rest of the columns into lists (of the same length)
  2. We derive the list "signal_0_mean" from performing ewm per element of the "signal_0" list
  3. We zip-like explode on the three lists at the same time
  4. Sort to recover order

The result:

shape: (10, 4)
┌────────┬─────────────────────┬───────────┬───────────────┐
│ symbol ┆ ts                  ┆ signal_0  ┆ signal_0_mean │
│ ---    ┆ ---                 ┆ ---       ┆ ---           │
│ cat    ┆ datetime[ns]        ┆ f64       ┆ f64           │
╞════════╪═════════════════════╪═══════════╪═══════════════╡
│ A      ┆ 2022-02-14 09:20:00 ┆ -1.704301 ┆ -1.704301     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 09:20:00 ┆ -1.181743 ┆ -1.181743     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 09:50:00 ┆ 1.040125  ┆ -0.2845       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 09:50:00 ┆ 0.776798  ┆ -0.168547     │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ...    ┆ ...                 ┆ ...       ┆ ...           │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 10:50:00 ┆ 2.073418  ┆ 0.940226      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 10:50:00 ┆ 1.623698  ┆ 0.752741      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A      ┆ 2022-02-14 11:20:00 ┆ 2.088835  ┆ 1.202843      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ AA     ┆ 2022-02-14 11:20:00 ┆ 1.741544  ┆ 0.9788        │
└────────┴─────────────────────┴───────────┴───────────────┘

Which exactly matches the expected result from pandas.

user_37811800
Related