Как вернуть несколько статистических данных в виде нескольких столбцов в контексте Polars grouby? ⇐ Python

[Anonymous]

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

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

import numpy as np
import polars as pl
import statsmodels.api as sm

from functools import partial

def ols_stats(s, yvar, xvars):
df = s.struct.unnest()
yvar = df[yvar].to_numpy()
xvars = df[xvars].to_numpy()
reg = sm.OLS(yvar, sm.add_constant(xvars), missing="drop").fit()
return np.concatenate((reg.params, reg.tvalues))

df = pl.DataFrame(
{
"day": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3],
"y": [1, 6, 3, 2, 8, 4, 5, 2, 7, 3, 1],
"x1": [1, 8, 2, 3, 5, 2, 1, 2, 7, 3, 1],
"x2": [8, 5, 3, 6, 3, 7, 3, 2, 9, 1, 1],
}
)

df.group_by("day").agg(
pl.struct("y", "x1", "x2")
.map_elements(partial(ols_stats, yvar="y", xvars=["x1", "x2"]))
.alias("params")
)

Результат приведенного выше фрагмента кода равен

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

shape: (3, 2)
┌─────┬─────────────────────────────────┐
│ day ┆ params                          │
│ --- ┆ ---                             │
│ i64 ┆ object                          │
╞═════╪═════════════════════════════════╡
│ 2   ┆ [2.0462002  0.22397054 0.33679… │
│ 1   ┆ [ 4.86623165  0.64029364 -0.65… │
│ 3   ┆ [0.5 0.5 0.  0. ]               │
└─────┴─────────────────────────────────┘

Как мне разделить «параметры» на отдельные столбцы с одним скалярным значением в каждом столбце?
Кроме того, мой код, похоже, дает сбой в некоторых крайних случаях. Ниже один из них.

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

df = pl.DataFrame(
{
"day": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3],
"y": [1, 6, 3, 2, 8, 4, 5, 2, 7, 3, None],
"x1": [1, 8, 2, 3, 5, 2, 1, 2, 7, 3, 1],
"x2": [8, 5, 3, 6, 3, 7, 3, 2, 9, 1, 1],
}
)

df.group_by("day").agg(
pl.struct("y", "x1", "x2")
.map_elements(partial(ols_stats, yvar="y", xvars=["x1", "x2"]))
.alias("params")
)

# ComputeError: ValueError: exog is not 1d or 2d

Как сделать код устойчивым к таким случаям?
Спасибо за помощь. И не стесняйтесь предлагать собственное решение.

Подробнее здесь: https://stackoverflow.com/questions/755 ... by-context