Add exposure kwarg to yli.regress, for Poisson etc.
This commit is contained in:
parent
17bf3cbcab
commit
fa95f6d75c
@ -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 == '<m8[ns]':
|
||||
model_kwargs['exposure'] = df[exposure].dt.total_seconds()
|
||||
else:
|
||||
model_kwargs['exposure'] = df[exposure]
|
||||
|
||||
# Fit model
|
||||
model = model_class(endog=dmatrices[0], exog=dmatrices[1], **model_kwargs)
|
||||
model.formula = dep + ' ~ ' + formula
|
||||
|
Loading…
Reference in New Issue
Block a user