Поляры: как обрабатывать «выражение окна не разрешено при агрегации»? ⇐ Python

[Anonymous]

У меня есть задача, в которой я хочу выполнить регрессию преобразованных столбцов для набора указанных столбцов в кадре данных Polars. Преобразование и набор независимых столбцов управляются спецификацией.
Ниже приведен упрощенный мини-пример для иллюстрации.

Код: Выделить всё

from functools import partial

import polars as pl
import numpy as np

def ols_fitted(s: pl.Series, yvar: str, xvars: list[str]) -> pl.Series:
df = s.struct.unnest()
y = df[yvar].to_numpy()
X = df[xvars].to_numpy()
fitted = np.dot(X, np.linalg.lstsq(X, y, rcond=None)[0])
return pl.Series(values=fitted, nan_to_null=True)

df = pl.DataFrame(
{
"date": [1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
"id": [1, 1, 1, 2, 2, 2, 2, 3, 3, 3],
"y": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
"g1": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
"g2": [1, 1, 1, 2, 2, 2, 3, 3, 3, 3],
"g3": [1, 1, 2, 2, 2, 3, 3, 4, 4, 4],
"x1": [2, 5, 4, 7, 3, 2, 5, 6, 7, 2],
"x2": [1, 5, 3, 4, 5, 6, 4, 3, 2, 1],
"x3": [3, 6, 8, 6, 4, 7, 5, 4, 8, 1],
}
)

specs = {
"first": {"yvar": "y", "gvars": ["g1"], "xvars": ["x1"]},
"second": {"yvar": "y", "gvars": ["g1", "g2"], "xvars": ["x1", "x2"]},
"third": {"yvar": "y", "gvars": ["g2", "g3"], "xvars": ["x2", "x3"]},
}

df.with_columns(
pl.struct(
(
pl.col(specs[specnm]["yvar"])
- pl.col(specs[specnm]["yvar"]).mean().over(specs[specnm]["gvars"])
).abs(),
*specs[specnm]["xvars"],
)
.map_elements(
partial(
ols_fitted, yvar=specs[specnm]["yvar"], xvars=specs[specnm]["xvars"]
)
)
.over("date", "id")
.alias(f"fitted_{specnm}")
for specnm in list(specs.keys())
)

Однако я получил следующую ошибку:

Код: Выделить всё

InvalidOperationError: window expression not allowed in aggregation

Не знаю, почему over не поддерживается в контексте агрегации. Было бы очень удобно, если бы это было так, как в моем примере.
Итак, мой реальный вопрос: как с этим справиться в моем конкретном случае? И, если с этим невозможно справиться, есть ли альтернативные способы заставить мой код работать систематически?

Подробнее здесь: https://stackoverflow.com/questions/756 ... ggregation