pythonpandasstatsmodelspatsy

How do I get the columns that a statsmodels / patsy formula depends on?


Suppose I have a pandas dataframe:

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]})

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()

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

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

Or even:

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

How can I get this from just the fit object?


Solution

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

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