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

假设我有一个熊猫数据框：

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.predict().例如,如果我尝试使用list(fit.params.index),则会得到：

['Intercept', 'x1:x2']

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

['x1', 'x2']

甚至：

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

我怎样才能从适合的对象得到这个？

解决方法:

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

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