[Python] How to return multiple stats as multiple columns in Polars grouby context?

The task at hand is to do multiple linear regression over multiple columns in groupby context and return respective beta coefficients and their associated t-values in separate columns.

Below is an illustration of an attempt to do such using statsmodels.

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")
)


The result from the code snippet above evaluates to

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. ] │
└─────┴─────────────────────────────────┘


How am I supposed to split the 'params' into separate columns with one scalar value in each column?

Also, my code seems to fail at some corner cases. Below is one of them.

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


How can I make the code robust to such case?

Thanks for your help. And feel free to suggest your own solution.

Continue reading...