如何用Numpy和Statsmodels实现将X列作为分类变量的OLS回归函数？

问题背景 #

给定一个接受numpy数组向量y和二维矩阵X的Python函数，原函数实现了普通OLS回归，但需要修改为将X的每一列都当作独立分类变量来进行回归，仅允许使用numpy和statsmodels库。原函数代码如下：

import statsmodels.api as sm

def f(y, X):
    return sm.OLS(y, sm.add_constant(X), missing="drop").fit()

修改方案 #

要将X的列作为分类变量处理，核心是把每个分类列转换为虚拟变量（one-hot编码），同时自动移除每个分类的基准类别以避免多重共线性。可以借助statsmodels内置的工具函数完成这个转换，修改后的函数如下：

import numpy as np
import statsmodels.api as sm

def ols_categorical(y, X):
    # 存储每一列转换后的虚拟变量矩阵
    dummy_list = []
    for col_idx in range(X.shape[1]):
        # 将当前列转为分类虚拟变量，drop=True移除基准类别
        dummy_col = sm.tools.tools.categorical(X[:, col_idx], drop=True)
        dummy_list.append(dummy_col)
    
    # 横向拼接所有虚拟变量矩阵
    X_dummies = np.hstack(dummy_list)
    # 添加常数项并拟合模型
    X_with_const = sm.add_constant(X_dummies, missing="drop")
    model = sm.OLS(y, X_with_const, missing="drop").fit()
    
    # 返回回归结果摘要
    return model.summary()

关键说明 #

• sm.tools.tools.categorical()：这个函数会自动将输入的分类列转换为虚拟变量矩阵，drop=True参数确保每个分类变量只保留n-1个虚拟变量（n为类别数），避免OLS回归中的多重共线性问题。
• 矩阵拼接：用np.hstack()将所有列的虚拟变量矩阵横向合并，得到最终的自变量矩阵。
• 结果返回：直接返回模型的摘要信息，包含系数、R²、显著性检验等完整回归结果。

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