diff --git a/yli/regress.py b/yli/regress.py index e770380..3965d20 100644 --- a/yli/regress.py +++ b/yli/regress.py @@ -689,7 +689,7 @@ def regress( model_class, df, dep, formula, *, nan_policy='warn', model_kwargs=None, fit_kwargs=None, - family=None, # common model_kwargs + family=None, exposure=None, # common model_kwargs cov_type=None, method=None, maxiter=None, start_params=None, # common fit_kwargs bool_baselevels=False, exp=None, _dmatrices=None, @@ -705,6 +705,8 @@ def regress( :type dep: str :param formula: Patsy formula for the regression model :type formula: str + :param exposure: Column in *df* for the exposure variable (numeric, some models only) + :type exposure: str :param nan_policy: How to handle *nan* values (see :ref:`nan-handling`) :type nan_policy: str :param model_kwargs: Keyword arguments to pass to *model_class* constructor @@ -789,7 +791,10 @@ def regress( if _dmatrices is None: # Check for/clean NaNs - df = df[[dep] + cols_for_formula(formula, df)] + if exposure is None: + df = df[[dep] + cols_for_formula(formula, df)] + else: + df = df[[dep, exposure] + cols_for_formula(formula, df)] df = check_nan(df, nan_policy) # Ensure numeric type for dependent variable @@ -808,6 +813,12 @@ def regress( # FIXME: Check before dropping dmatrices = (dmatrices[0], dmatrices[1].iloc[:,1:]) + if exposure is not None: + if df[exposure].dtype == '