From fa95f6d75c69f921f2c304d347fe2bb7f03e696a Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Sat, 25 Feb 2023 14:12:44 +1100
Subject: [PATCH] Add exposure kwarg to yli.regress, for Poisson etc.

---
 yli/regress.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

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