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

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

Answer 2

由于您所要求的只是简单的线性回归残差，我们只需几个 Polars 表达式即可做到这一点：

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_columns(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  │
└───────────┴──────────┴───────────┘

Answer 3

晚了一年！但这也与我自己相关，因此最近发布了一个 Polars 扩展包，用于解决这些问题polars-ols。

pip install polars-ols

具有多个特征的最小二乘就简单地变成了极坐标表达式，所以在你的例子中是这样的：

import polars as pl
import polars_ols as pls  # registers .least_squares namespace

df = df.with_columns(
    pl.col("y").least_squares.ols(pl.col("x1"), pl.col("x2"), add_intercept=True, mode="residuals")
      .over("day")
      .alias("residuals")
)

它可以在 Rust 中完成所有操作，就像其他 Polars 表达式一样，因此也可以与 .over() 很好地并行。分享以防仍然有帮助！