如何获取statsmodels/patsy公式所依赖的列? [英] How do I get the columns that a statsmodels / patsy formula depends on?
问题描述
假设我有一个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']
这篇关于如何获取statsmodels/patsy公式所依赖的列?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!