Сгруппированное скользящее среднее в полярах ⇐ Python

[Anonymous]

Здесь задан аналогичный вопрос
Однако в моем случае это не сработало.
У меня есть фрейм данных с тремя столбцами, датой, группами и проблемой. Я хочу создать трехдневное скользящее среднее значений столбца проблем, сгруппированных по группам и дате. Однако, следуя приведенному выше ответу, я получил все нули.
import polars as pl
from datetime import date
import numpy as np

dates = pl.date_range(date(2024, 12, 1), date(2024, 12, 30), "1d", eager=True).alias("date")

days = pl.concat([dates,dates])
groups = pl.concat([pl.select(pl.repeat("B", n = 30)).to_series(),
pl.select(pl.repeat("A", n = 30)).to_series()]).alias('groups')

data = pl.DataFrame([days, groups])

data2 = data.with_columns(pl.lit(np.random.rand(data.height)).alias("prob"))
data2.with_columns(
pl.col('prob')
.rolling_mean(window_size = 3)
.over('date','groups')
.alias('rolling_mean')
)

shape: (60, 4)
┌────────────┬────────┬──────────┬──────────────┐
│ date ┆ groups ┆ prob ┆ rolling_mean │
│ --- ┆ --- ┆ --- ┆ --- │
│ date ┆ str ┆ f64 ┆ f64 │
╞════════════╪════════╪══════════╪══════════════╡
│ 2024-12-01 ┆ B ┆ 0.938982 ┆ null │
│ 2024-12-02 ┆ B ┆ 0.103133 ┆ null │
│ 2024-12-03 ┆ B ┆ 0.724672 ┆ null │
│ 2024-12-04 ┆ B ┆ 0.495868 ┆ null │
│ 2024-12-05 ┆ B ┆ 0.621124 ┆ null │
│ … ┆ … ┆ … ┆ … │
│ 2024-12-26 ┆ A ┆ 0.762529 ┆ null │
│ 2024-12-27 ┆ A ┆ 0.766366 ┆ null │
│ 2024-12-28 ┆ A ┆ 0.272936 ┆ null │
│ 2024-12-29 ┆ A ┆ 0.28709 ┆ null │
│ 2024-12-30 ┆ A ┆ 0.403478 ┆ null │
└────────────┴────────┴──────────┴──────────────┘

В документации я нашел .rolling_mean_by и попробовал использовать его вместо этого, но вместо скользящего среднего он, похоже, просто возвращает пробное значение для каждой строки.
data2.with_columns(
pl.col('prob')
.rolling_mean_by(window_size = '3d', by = 'date')
.over('groups', 'date')
.alias('rolling_mean')
)

shape: (60, 4)
┌────────────┬────────┬──────────┬──────────────┐
│ date ┆ groups ┆ prob ┆ rolling_mean │
│ --- ┆ --- ┆ --- ┆ --- │
│ date ┆ str ┆ f64 ┆ f64 │
╞════════════╪════════╪══════════╪══════════════╡
│ 2024-12-01 ┆ B ┆ 0.938982 ┆ 0.938982 │
│ 2024-12-02 ┆ B ┆ 0.103133 ┆ 0.103133 │
│ 2024-12-03 ┆ B ┆ 0.724672 ┆ 0.724672 │
│ 2024-12-04 ┆ B ┆ 0.495868 ┆ 0.495868 │
│ 2024-12-05 ┆ B ┆ 0.621124 ┆ 0.621124 │
│ … ┆ … ┆ … ┆ … │
│ 2024-12-26 ┆ A ┆ 0.762529 ┆ 0.762529 │
│ 2024-12-27 ┆ A ┆ 0.766366 ┆ 0.766366 │
│ 2024-12-28 ┆ A ┆ 0.272936 ┆ 0.272936 │
│ 2024-12-29 ┆ A ┆ 0.28709 ┆ 0.28709 │
│ 2024-12-30 ┆ A ┆ 0.403478 ┆ 0.403478 │
└────────────┴────────┴──────────┴──────────────┘


Подробнее здесь: https://stackoverflow.com/questions/794 ... -in-polars