improve manual type checking sections in the experiment modules

Author: drbenvincent

causalpy/expt_diff_in_diff.py

@@ -19,6 +19,7 @@
 from causalpy.data_validation import DiDDataValidator
 from causalpy.experiments import ExperimentalDesign
 from causalpy.pymc_models import PyMCModel
+from causalpy.skl_models import ScikitLearnModel
 from causalpy.utils import convert_to_string
 
 
@@ -85,13 +86,14 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # ******** THIS IS SUBOPTIMAL AT THE MOMENT ************************************
+        # fit model
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.X.shape[0])}
             self.model.fit(X=self.X, y=self.y, coords=COORDS)
-        else:
+        elif isinstance(self.model, ScikitLearnModel):
             self.model.fit(X=self.X, y=self.y)
-        # ******************************************************************************
+        else:
+            raise ValueError("Model type not recognized")
 
         # predicted outcome for control group
         self.x_pred_control = (
@@ -151,25 +153,23 @@ def __init__(
                 new_x.iloc[:, i] = 0
         self.y_pred_counterfactual = self.model.predict(np.asarray(new_x))
 
-        # ******** THIS IS SUBOPTIMAL AT THE MOMENT ************************************
+        # calculate causal impact
         if isinstance(self.model, PyMCModel):
-            # calculate causal impact &&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
             # This is the coefficient on the interaction term
             coeff_names = self.model.idata.posterior.coords["coeffs"].data
             for i, label in enumerate(coeff_names):
                 if "post_treatment" in label and self.group_variable_name in label:
                     self.causal_impact = self.model.idata.posterior["beta"].isel(
                         {"coeffs": i}
                     )
-            # &&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
-        else:
-            # calculate causal impact
+        elif isinstance(self.model, ScikitLearnModel):
             # This is the coefficient on the interaction term
-            # TODO: THIS IS NOT YET CORRECT
+            # TODO: THIS IS NOT YET CORRECT ?????
             self.causal_impact = (
                 self.y_pred_treatment[1] - self.y_pred_counterfactual[0]
             )[0]
-        # ******************************************************************************
+        else:
+            raise ValueError("Model type not recognized")
 
     def plot(self, round_to=None):
         """

causalpy/expt_prepostfit.py

@@ -20,6 +20,7 @@
 from causalpy.data_validation import PrePostFitDataValidator
 from causalpy.experiments import ExperimentalDesign
 from causalpy.pymc_models import PyMCModel
+from causalpy.skl_models import ScikitLearnModel
 
 
 class PrePostFit(ExperimentalDesign, PrePostFitDataValidator):
@@ -62,13 +63,13 @@ def __init__(
         self.post_y = np.asarray(new_y)
 
         # fit the model to the observed (pre-intervention) data
-        # ******** THIS IS SUBOPTIMAL AT THE MOMENT ************************************
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.pre_X.shape[0])}
             self.model.fit(X=self.pre_X, y=self.pre_y, coords=COORDS)
-        else:
+        elif isinstance(self.model, ScikitLearnModel):
             self.model.fit(X=self.pre_X, y=self.pre_y)
-        # ******************************************************************************
+        else:
+            raise ValueError("Model type not recognized")
 
         # score the goodness of fit to the pre-intervention data
         self.score = self.model.score(X=self.pre_X, y=self.pre_y)

causalpy/expt_prepostnegd.py

@@ -19,6 +19,7 @@
 from causalpy.data_validation import PrePostNEGDDataValidator
 from causalpy.experiments import ExperimentalDesign
 from causalpy.pymc_models import PyMCModel
+from causalpy.skl_models import ScikitLearnModel
 from causalpy.utils import round_num
 
 
@@ -94,13 +95,13 @@ def __init__(
         self.outcome_variable_name = y.design_info.column_names[0]
 
         # fit the model to the observed (pre-intervention) data
-        # ******** THIS IS SUBOPTIMAL AT THE MOMENT ************************************
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.X.shape[0])}
             self.model.fit(X=self.X, y=self.y, coords=COORDS)
-        else:
+        elif isinstance(self.model, ScikitLearnModel):
             raise NotImplementedError("Not implemented for OLS model")
-        # ******************************************************************************
+        else:
+            raise ValueError("Model type not recognized")
 
         # Calculate the posterior predictive for the treatment and control for an
         # interpolated set of pretest values

causalpy/expt_regression_discontinuity.py

@@ -21,6 +21,7 @@
 from causalpy.data_validation import RDDataValidator
 from causalpy.experiments import ExperimentalDesign
 from causalpy.pymc_models import PyMCModel
+from causalpy.skl_models import ScikitLearnModel
 from causalpy.utils import convert_to_string
 
 
@@ -105,14 +106,15 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # ******** THIS IS SUBOPTIMAL AT THE MOMENT ************************************
+        # fit model
         if isinstance(self.model, PyMCModel):
             # fit the model to the observed (pre-intervention) data
             COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.X.shape[0])}
             self.model.fit(X=self.X, y=self.y, coords=COORDS)
-        else:
+        elif isinstance(self.model, ScikitLearnModel):
             self.model.fit(X=self.X, y=self.y)
-        # ******************************************************************************
+        else:
+            raise ValueError("Model type not recognized")
 
         # score the goodness of fit to all data
         self.score = self.model.score(X=self.X, y=self.y)