Polars是否提供最小二乘线性回归API？仅用Polars如何实现？

关于Polars的线性回归API #

Polars作为专注于高性能数据处理的库，没有内置的最小二乘线性回归API，这类统计建模功能并非它的核心定位，所以在官方API文档中找不到相关接口。

仅用Polars实现最小二乘线性回归 #

我们可以基于最小二乘的数学原理，利用Polars的线性代数工具和表达式系统高效实现。最小二乘的系数计算公式为：
$$\beta = (X^T X)^{-1} X^T y$$
其中$X$是包含截距项的特征矩阵，$y$是目标变量向量。

实现步骤与代码 #

1. 为数据添加截距列（全1的列）
2. 提取特征矩阵$X$和目标向量$y$
3. 利用Polars的线性代数函数计算系数

结合测试数据，完整实现代码如下：

import polars as pl

data = {
    'x': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
    'y': [22.0, 33.9, 44.8, 78.9, 44.3, 20.5, 30.5, 56.4, 92.3, 22.1, 88, 10.1]
}

df = pl.DataFrame(data)

# 添加截距列
df_with_intercept = df.with_columns(pl.lit(1).alias('intercept'))

# 提取X和y的矩阵形式（Polars DataFrame）
X = df_with_intercept.select('intercept', 'x')
y = df_with_intercept.select('y')

# 计算X^T X
X_T_X = pl.linalg.dot(X.transpose(), X)
# 计算(X^T X)的逆
X_T_X_inv = pl.linalg.inv(X_T_X)
# 计算X^T y
X_T_y = pl.linalg.dot(X.transpose(), y)
# 计算系数β
beta = pl.linalg.dot(X_T_X_inv, X_T_y)

# 输出系数：第一个是截距，第二个是x的系数
print("回归系数：")
print(beta)

# 可选：用系数计算预测值
df_pred = df_with_intercept.with_columns(
    (pl.col('intercept') * beta[0,0] + pl.col('x') * beta[1,0]).alias('y_pred')
)
print("\n带预测值的数据集：")
print(df_pred)

说明 #

• 全程使用Polars的pl.linalg模块中的线性代数函数，无需切换到其他库，保证了实现的纯粹性
• 针对大规模数据集，Polars的并行优化特性能维持高效的计算速度

内容的提问来源于stack exchange，提问作者Young Eddie