Скольжение наклона линейной регрессии по n строкам по группам ⇐ Python

[Anonymous]

Мои данные указаны на уровне ID-месяца с оплатой каждый месяц. Он сортируется по идентификатору и dt. Я хотел бы создать новый столбец, который для каждой группы будет содержать линейный наклон для следующих N месяцев платежей. Вот образец:
import polars as pl

data = {"id": ['a','a','a','a','a','a'], "dt": ['2024-01-01', '2024-02-01', '2024-03-01', '2024-04-01', '2024-05-01', '2024-06-01',], "pmt": [3341,3205,3287,3544,6536,5994]}
df = pl.DataFrame(data, schema={"id": pl.String, "dt": pl.Date, 'pmt':pl.Int64}).with_columns(pl.col("dt").set_sorted())
print(df)

Это выглядит следующим образом:
shape: (6, 3)
┌─────┬────────────┬──────┐
│ id ┆ dt ┆ pmt │
│ --- ┆ --- ┆ --- │
│ str ┆ date ┆ i64 │
╞═════╪════════════╪══════╡
│ a ┆ 2024-01-01 ┆ 3341 │
│ a ┆ 2024-02-01 ┆ 3205 │
│ a ┆ 2024-03-01 ┆ 3287 │
│ a ┆ 2024-04-01 ┆ 3544 │
│ a ┆ 2024-05-01 ┆ 6536 │
│ a ┆ 2024-06-01 ┆ 5994 │
└─────┴────────────┴──────┘

Я знаю, что наклон линейной регрессии для этого идентификатора для этого шестимесячного блока равен 671,86, поэтому это значение будет находиться в первой строке. В следующей строке будет указан наклон для следующих 6 месяцев и т. д. Результат будет выглядеть так (очевидно, что строки 2–6 содержат поддельные данные наклона; просто для наглядности):
shape: (6, 4)
┌─────┬────────────┬──────┬───────────────────┐
│ id ┆ dt ┆ pmt ┆ slope_pmt_next6mo │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ date ┆ i64 ┆ f32 │
╞═════╪════════════╪══════╪═══════════════════╡
│ a ┆ 2024-01-01 ┆ 3341 ┆ 671.859985 │
│ a ┆ 2024-02-01 ┆ 3205 ┆ 700.25 │
│ a ┆ 2024-03-01 ┆ 3287 ┆ 646.210022 │
│ a ┆ 2024-04-01 ┆ 3544 ┆ 683.880005 │
│ a ┆ 2024-05-01 ┆ 6536 ┆ 547.179993 │
│ a ┆ 2024-06-01 ┆ 5994 ┆ 525.48999 │
└─────┴────────────┴──────┴───────────────────┘

Я предполагаю, что я бы использовал вариант Polars.DataFrame.rolling, но я не совсем понимаю синтаксис. Я попробовал следующее, но в результате в первых двух строках создается NaN, поэтому не знаю, что происходит:
def ols_slope(y: pl.Expr) -> pl.Expr:
# Calculate linear regression slope
x = y.rank("ordinal")
numerator = ((x - x.mean())*(y - y.mean())).sum()
denominator = ((x - x.mean())**2).sum()
return numerator / denominator

(
df
.rolling(index_column=("dt"), period="6mo", closed='none')
.agg(ols_slope(pl.col("pmt")).alias("pmt_slope"))
)

shape: (6, 2)
┌────────────┬───────────┐
│ dt ┆ pmt_slope │
│ --- ┆ --- │
│ date ┆ f64 │
╞════════════╪═══════════╡
│ 2024-01-01 ┆ NaN │
│ 2024-02-01 ┆ NaN │
│ 2024-03-01 ┆ 136.0 │
│ 2024-04-01 ┆ 68.0 │
│ 2024-05-01 ┆ 107.1 │
│ 2024-06-01 ┆ 691.9 │
└────────────┴───────────┘


Подробнее здесь: https://stackoverflow.com/questions/789 ... s-by-group