如何在极坐标选择或分组上下文中进行回归（例如简单线性）？

Question

我用极地代替熊猫。我对速度和惰性计算/评估感到非常惊讶。目前，有很多关于惰性数据帧的方法，但它们只能驱使我到目前为止。

因此，我想知道将极坐标与其他工具结合使用以实现更复杂的操作（例如回归/模型拟合）的最佳方法是什么。

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

假设我有一个包含 day、y、x1 和 x2 列的极坐标数据框，并且我想生成一个序列，它是按天对 x1 和 x2 进行回归 y 的残差。我包含了如下代码示例以及如何使用 pandas 和 statsmodels 解决它。如何使用惯用的极坐标以最有效的方式获得相同的结果？

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。

由于 Polars 不使用numpy内存statsmodels，因此您必须转换 Polars 类型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.unnest()
    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.group_by("day")
   .agg(
       pl.struct(["y", "x1", "x2"]).map_elements(partial(regress_resid, yvar="y", xvars=["x1", "x2"]))
   )
)