如何做回归(简单的线性例如)在极地选择或groupby上下文?

Question

我用极地代替熊猫。我对速度和懒惰的计算/评估感到非常惊讶。现在，有很多关于lazy dataframe的方法，但是它们只能驱动我到此为止。

所以，我想知道将极性与其他工具结合使用以实现更复杂的操作(如回归/模型拟合)的最佳方法是什么。

更具体地说，我将给出一个涉及线性回归的例子。

假设我有一个包含day, y, x1和x2列的极数据框架，我想生成一个序列，它是y在x1和x2组上按day回归的残差。我包含了如下代码示例，以及如何使用pandas和statmodel来解决这个问题。如何使用惯用极性以最高效的方式获得相同的结果?

import pandas as pd
import statsmodels.api as sm
def regress_resid(df, yvar, xvars):
result = sm.OLS(df[yvar], sm.add_constant(df[xvars])).fit()
return result.resid
df = pd.DataFrame(
{
"day": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
"y": [1, 6, 3, 2, 8, 4, 5, 2, 7, 3],
"x1": [1, 8, 2, 3, 5, 2, 1, 2, 7, 3],
"x2": [8, 5, 3, 6, 3, 7, 3, 2, 9, 1],
}
)
df.groupby("day").apply(regress_resid, "y", ["x1, "x2])
# day
# 1    0    0.772431
#      1   -0.689233
#      2   -1.167210
#      3   -0.827896
#      4    1.911909
# 2    5   -0.851691
#      6    1.719451
#      7   -1.167727
#      8    0.354871
#      9   -0.054905

谢谢你的帮助。

Answer 1

如果要将多个列传递给函数，则必须将它们打包到Struct中，因为极性表达式总是从Series -> Series映射。

因为极性不像statsmodels那样使用numpy内存，所以必须转换极性类型to_numpy。对于一维结构，这通常是免费的。

最后，该函数不应该返回numpy数组，而是一个极地Series，所以我们转换结果。

import polars as pl
from functools import partial
import statsmodels.api as sm
def regress_resid(s: pl.Series, yvar: str, xvars: list[str]) -> pl.Series:
df = s.struct.to_frame()
yvar = df[yvar].to_numpy()
xvars = df[xvars].to_numpy()

result = sm.OLS(yvar, sm.add_constant(xvars)).fit()
return pl.Series(result.resid)

df = pl.DataFrame(
{
"day": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
"y": [1, 6, 3, 2, 8, 4, 5, 2, 7, 3],
"x1": [1, 8, 2, 3, 5, 2, 1, 2, 7, 3],
"x2": [8, 5, 3, 6, 3, 7, 3, 2, 9, 1],
}
)
(df.groupby("day")
.agg([
pl.struct(["y", "x1", "x2"]).apply(partial(regress_resid, yvar="y", xvars=["x1", "x2"]))
])
)

Answer 2

因为所有你要求的是简单的线性回归残差，我们可以做到这一点，只需几个极地表达式:

import numpy as np
import polars as pl
# create some correlated data
generator = np.random.default_rng(seed=4)
x = generator.normal(size=100)
error = generator.normal(size=100)
df = pl.DataFrame({"x": x, "y": 2 + x + error})
def residuals(x: pl.Expr, y: pl.Expr) -> pl.Expr:
# e_i = y_i - a - bx_i
#     = y_i - ȳ + bx̄ - bx_i
#     = y_i - ȳ - b(x_i - x̄)
x_mean = x.mean()
y_mean = y.mean()
x_demeaned = x - x.mean()
y_demeaned = y - y.mean()
x_demeaned_squared = x_demeaned.pow(2)
beta = x_demeaned.dot(y_demeaned) / x_demeaned_squared.sum()
return y_demeaned - beta * x_demeaned
print(
df
.with_column(residuals(pl.col("x"), pl.col("y")).alias("e"))
.head(3) # remove this to get the entire result
)

前几行如下所示:

shape: (3, 3)
┌───────────┬──────────┬───────────┐
│ x         ┆ y        ┆ e         │
│ ---       ┆ ---      ┆ ---       │
│ f64       ┆ f64      ┆ f64       │
╞═══════════╪══════════╪═══════════╡
│ -0.651791 ┆ 1.257485 ┆ -0.146135 │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ -0.174717 ┆ 1.704334 ┆ -0.260438 │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 1.663724  ┆ 4.740446 ┆ 0.613234  │
└───────────┴──────────┴───────────┘