Uma pergunta semelhante é feita aqui
No entanto, não pareceu funcionar no meu caso.
Tenho um dataframe com 3 colunas, data, grupos, prob. O que quero é criar uma média móvel de 3 dias dos valores da coluna prob agrupados por grupos e data. No entanto, seguindo a resposta vinculada acima, obtive todos os nulos retornados.
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")
len(dates)
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(
rolling_mean =
pl.col('prob')
.rolling_mean(window_size = 3)
.over('date','groups')
)
"""
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 │
└────────────┴────────┴──────────┴──────────────┘
""""
Na documentação, encontrei .rolling_mean_by e tentei usá-lo, mas em vez de fazer uma média móvel, ele parece apenas retornar o valor prob para cada linha.
data2.with_columns(
rolling_mean =
pl.col('prob')
.rolling_mean_by(window_size = '3d', by = 'date')
.over('groups', 'date')
)
"""
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 │
└────────────┴────────┴──────────┴──────────────┘
""""
