pylint

Author: jmafoster1

causal_testing/estimation/abstract_regression_estimator.py

@@ -5,7 +5,7 @@
 from typing import Any
 
 import pandas as pd
-from patsy import dmatrix, dmatrices  # pylint: disable = no-name-in-module
+from patsy import dmatrices, dmatrix  # pylint: disable = no-name-in-module
 from statsmodels.regression.linear_model import RegressionResultsWrapper
 
 from causal_testing.estimation.abstract_estimator import Estimator
@@ -113,6 +113,20 @@ def fit_model(self, data=None) -> RegressionResultsWrapper:
             model = self.regressor(data[self.base_test_case.outcome_variable.name], data[self.covariates]).fit(disp=0)
         return model
 
+    def treatment_columns(self, model: RegressionResultsWrapper):
+        """
+        Get the names of the treatment columns from the model.
+        This is a workaround for statsmodels mangling the names of categorical variables to include the values.
+
+        :param model: The fitted model from which to extract the variable names.
+        """
+        return [
+            param
+            for param in model.params.index
+            if param == self.base_test_case.treatment_variable.name
+            or param.startswith(self.base_test_case.treatment_variable.name + "[")
+        ]
+
     def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
         """Estimate the outcomes under control and treatment.
 

causal_testing/estimation/logistic_regression_estimator.py

@@ -41,13 +41,7 @@ def estimate_unit_odds_ratio(self) -> EffectEstimate:
         """
         model = self.fit_model(self.df)
 
-        treatment_columns = [
-            param
-            for param in model.params.index
-            if param == self.base_test_case.treatment_variable.name
-            or param.startswith(self.base_test_case.treatment_variable.name + "[")
-        ]
-
+        treatment_columns = self.treatment_columns(model)
         confidence_intervals = np.exp(model.conf_int(self.alpha).loc[treatment_columns])
 
         result = EffectEstimate(

causal_testing/estimation/multinomial_regression_estimator.py

@@ -36,16 +36,9 @@ def estimate_unit_odds_ratio(self) -> EffectEstimate:
         """
         model = self.fit_model(self.df)
 
-        treatment_columns = [
-            param
-            for param in model.params.index
-            if param == self.base_test_case.treatment_variable.name
-            or param.startswith(self.base_test_case.treatment_variable.name + "[")
-        ]
-
         conf_int = model.conf_int(self.alpha)
         levels_of_interest = [
-            (level, covariate) for level, covariate in conf_int.index if covariate in treatment_columns
+            (level, covariate) for level, covariate in conf_int.index if covariate in self.treatment_columns(model)
         ]
         confidence_intervals = np.exp(conf_int.loc[levels_of_interest])