add input validation and associated tests for exceptions for RegressionDiscontinuity

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -560,16 +560,15 @@ def __init__(
         self.formula = formula
         self.running_variable_name = running_variable_name
         self.treatment_threshold = treatment_threshold
+        self._input_validation()
+
         y, X = dmatrices(formula, self.data)
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
         self.labels = X.design_info.column_names
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # TODO: `treated` is a deterministic function of x and treatment_threshold, so
-        # this could be a function rather than supplied data
-
         # DEVIATION FROM SKL EXPERIMENT CODE =============================
         # fit the model to the observed (pre-intervention) data
         COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.X.shape[0])}
@@ -610,6 +609,18 @@ def __init__(
             - self.pred_discon["posterior_predictive"].sel(obs_ind=0)["mu"]
         )
 
+    def _input_validation(self):
+        """Validate the input data and model formula for correctness"""
+        if "treated" not in self.formula:
+            raise FormulaException(
+                "A predictor called `treated` should be in the formula"
+            )
+
+        if _is_variable_dummy_coded(self.data["treated"]) is False:
+            raise DataException(
+                """The treated variable should be dummy coded. Consisting of 0's and 1's only."""  # noqa: E501
+            )
+
     def _is_treated(self, x):
         """Returns ``True`` if `x` is greater than or equal to the treatment threshold.
 

causalpy/tests/test_input_validation.py

@@ -161,3 +161,44 @@ def test_ancova_validation_2_levels():
             pretreatment_variable_name="pre",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         )
+
+
+# Regression discontinuity
+
+
+def test_rd_validation_treated_in_formula():
+    """Test that we get a FormulaException if treated is not in the model formula"""
+    df = pd.DataFrame(
+        {
+            "x": [0, 1, 2, 3],
+            "treated": [0, 0, 1, 1],
+            "y": [1, 1, 2, 2],
+        }
+    )
+
+    with pytest.raises(FormulaException):
+        _ = cp.pymc_experiments.RegressionDiscontinuity(
+            df,
+            formula="y ~ 1 + x",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            treatment_threshold=0.5,
+        )
+
+
+def test_rd_validation_treated_is_dummy():
+    """Test that we get a DataException if treated is not dummy coded"""
+    df = pd.DataFrame(
+        {
+            "x": [0, 1, 2, 3],
+            "treated": ["control", "control", "treated", "treated"],
+            "y": [1, 1, 2, 2],
+        }
+    )
+
+    with pytest.raises(DataException):
+        _ = cp.pymc_experiments.RegressionDiscontinuity(
+            df,
+            formula="y ~ 1 + x + treated",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            treatment_threshold=0.5,
+        )