RunasSudo
/
scipy-yli
1
Fork 0
Code Issues 1 Pull Requests Releases Activity

Add exposure kwarg to yli.regress, for Poisson etc.

This commit is contained in:
RunasSudo 2023-02-25 14:12:44 +11:00
parent 17bf3cbcab
commit fa95f6d75c
Signed by: RunasSudo
GPG Key ID: 7234E476BF21C61A
1 changed files with 13 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
yli/regress.py
View File

@ -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