如何获取statsmodels / patsy公式所依赖的列?

时间:2017-04-12 19:26:17

标签: python pandas statsmodels patsy

假设我有一个pandas数据帧:

df = pd.DataFrame({'x1': [0, 1, 2, 3, 4], 
                   'x2': [10, 9, 8, 7, 6], 
                   'x3': [.1, .1, .2, 4, 8], 
                   'y': [17, 18, 19, 20, 21]})

现在我使用公式(在引擎盖下使用statsmodels)拟合patsy模型:

import statsmodels.formula.api as smf
fit = smf.ols(formula='y ~ x1:x2', data=df).fit()

我想要的是df所依赖的fit列的列表,以便我可以在另一个数据集上使用fit.predict()。例如,如果我尝试list(fit.params.index),我会得到:

['Intercept', 'x1:x2']

我尝试重新创建patsy设计矩阵,并使用design_info,但我仍然只获得x1:x2。我想要的是:

['x1', 'x2']

甚至:

['Intercept', 'x1', 'x2']

如何从fit对象获取此内容?

3 个答案:

答案 0 :(得分:3)

只需测试列名称是否出现在公式的字符串表示中:

ols = smf.ols(formula='y ~ x1:x2', data=df)
fit = ols.fit()

print([c for c in df.columns if c in ols.formula])
['x1', 'x2', 'y']

通过重建patsy模型还有另一种方法(更详细,但也更可靠),并且它不依赖于原始数据框:

md = patsy.ModelDesc.from_formula(ols.formula)
termlist = md.rhs_termlist + md.lhs_termlist

factors = []
for term in termlist:
    for factor in term.factors:
        factors.append(factor.name())

print(factors)
['x1', 'x2', 'y']

答案 1 :(得分:0)

predict采用相同的数据框架或字典结构,并且调用patsy以兼容的方式转换它。要复制此内容,您还可以检查statsmodels.base.model.Results.predict核心

中的代码
exog = dmatrix(self.model.data.design_info.builder,
                           exog, return_type="dataframe")

公式信息本身存储在termsdesign_info的说明中。变量名称本身在summary()中使用,并在返回的pandas系列中用作索引,例如在results.params中。

答案 2 :(得分:0)

ols.exog_namesols.endog_names应该这样做