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

本文介绍了如何获取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]})

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

Now I fit a statsmodels model using a formula (which uses patsy under the hood):

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

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

What I want is a list of the columns of df that fit depends on, so that I can use fit.predict() on another dataset. If I try list(fit.params.index), for example, I get:

['Intercept', 'x1:x2']

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

I've tried recreating the patsy design matrix, and using design_info, but I still only ever get x1:x2. What I want is:

['x1', 'x2']

甚至:

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

如何仅从fit对象获得此信息?

How can I get this from just the fit object?

推荐答案

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

Simply test if the column names appear in the string representation of the formula:

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模型(更详细，但也更可靠)，并且它不依赖于原始数据帧:

There is another approach by reconstructing the patsy model (more verbose, but also more reliable) and it does not depend on the original data frame:

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