复制 Scipy 的 RegressionResults.predict 功能

Replicate Scipy's RegressionResults.predict functionality

这是我的示例程序:

import numpy as np
import pandas as pd
import statsmodels
from statsmodels.formula.api import ols
df = pd.DataFrame({"z": [1,1,1,2,2,2,3,3,3],
                   "x":[0,1,2,0,1,2,0,1,2],
                   "y":[0,2,4,3,5,7,7,9,11]
                   })
model = ols("y ~ x + z + I(z**2)", df).fit()
model.params

newdf = pd.DataFrame({"z": [4,4,4,5,5,5],
                   "x":[0,1,2,0,1,2]
                   })
model.predict(newdf)

你会注意到,如果你 运行 这个,model.params 是一个 pandas 系列,其索引与公式的 right-hand 边相同,除了附加条目:"Intercept"

>  Out[2]: 
>     Intercept   -2.0
>     x            2.0
>     z            1.5
>     I(z ** 2)    0.5
>     dtype: float64

并且,使用一些我无法确定的内部功能,RegressionResults object 的 .predict() 可以从 newdf 中识别列 headers,将它们匹配起来(包括 patsy 语法"I(z**2)"),添加截距,return 一个答案系列。 (这是我示例代码的最后一行)

这看起来很方便!比每当我想评估它的细微变化时在 python/numpy 代码中再次写出我的公式要好。我觉得应该有一些方法可以让我为公式系数构建一个类似的 pd.Series,而不是通过模型和拟合来创建它。然后我应该能够将其应用于适当的数据框,作为评估函数的一种方式。

我试图弄清楚 statsmodel 是如何做到这一点的,但没有成功,我在 patsy 的相关函数文档页面中没有发现任何明显的东西,我似乎也无法进入源代码的这一部分,而调试。 有人知道如何设置吗?

我最终拼凑出一种方法。

def predict(coeffs,datadf:pd.DataFrame)->np.array:
    """Apply a series (or df) of coefficents indexed by model terms to new data

    :param coeffs: a series whose elements are coefficients and index are the formula terms
                or a df whose column names are formula terms, and each row is a set of coefficients
    :param datadf: the new data to predict on
    """
    desc = patsy.ModelDesc([],[patsy.Term([]) if column=="Intercept" else patsy.Term([patsy.EvalFactor(column)]) for column in coeffs.index] )

    dmat = patsy.dmatrix(desc,datadf)
    return np.dot(dmat, coeffs.T)

newdf["y"] = predict(model.params,newdf)

如果有人感到困惑,这对我如此有吸引力的原因是我使用 df.groupby("column").apply(FitFunction) 分段拟合数据。 FitFunction() return model.params 系列似乎是 pandas 范式中最干净的方法。