如何获取 statsmodels / patsy 公式所依赖的列？答案

【问题标题】：How do I get the columns that a statsmodels / patsy formula depends on?如何获取 statsmodels / patsy 公式所依赖的列？
【发布时间】：2017-09-08 17:39:36
【问题描述】：

假设我有一个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()

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

['Intercept', 'x1:x2']

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

['x1', 'x2']

甚至：

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

我怎样才能从 fit 对象中得到这个？

【问题讨论】：

为什么不直接将'x1:x2' 拆分为':'，那么，如果您只是在交互x1 和x2？像fit.model.formula.split(':') 这样的东西，然后适当地过滤掉其余的东西。地狱，正则表达式拆分会更好，处理+、:等。
@bwk 你在这个问题上有什么进展吗？看看我的回答，应该符合你的需求。

标签： python pandas statsmodels patsy

【解决方案1】：

ols.exog_names 和 ols.endog_names 应该这样做

【讨论】：

【解决方案2】：

简单测试列名是否出现在公式的字符串表示中：

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']

【讨论】：

谢谢，这就是我要找的！
如果您碰巧有一个列名是公式的子字符串，这将不起作用。例如，我有一个名为 rt 的列，它由公式 C(other_col, Helmert) 捕获
@sammosummo 确实，我在答案中提到的第一种方法将匹配任何子字符串。但是，第二种方法不执行这种字符串匹配。
第二种方法没有回答 OP 的问题，因为它找到了 patsy 术语，而不是原始数据框中的列。使用分类编码因子时，这些是不一样的。

【解决方案3】：

predict 采用相同的数据帧或字典结构，调用 patsy 以兼容的方式对其进行转换。要复制这一点，您还可以检查statsmodels.base.model.Results.predict 中的代码，其核心是

exog = dmatrix(self.model.data.design_info.builder,
                           exog, return_type="dataframe")

公式信息本身存储在design_info中的terms的描述中。变量名称本身在summary() 中使用，并作为返回的pandas Series 中的索引，例如在results.params 中。

【讨论】：

获取公式不是问题。 summary()中的变量名不是DataFrame中的原始变量，而是变换后的变量名；即分别有“x1:x2”而不是“x1”和“x2”。