[IV 212] adding an input validation test

Author: NathanielF

causalpy/pymc_experiments.py

@@ -899,6 +899,7 @@ def __init__(
         self.formula = formula
         self.instruments_formula = instruments_formula
         self.model = model
+        self._input_validation()
 
         y, X = dmatrices(formula, self.data)
         self._y_design_info = y.design_info
@@ -954,3 +955,17 @@ def get_naive_OLS_fit(self):
         beta_params.insert(0, ols_reg.intercept_[0])
         self.ols_beta_params = dict(zip(self._x_design_info.column_names, beta_params))
         self.ols_reg = ols_reg
+
+    def _input_validation(self):
+        """Validate the input data and model formula for correctness"""
+        treatment = self.instruments_formula.split("~")[0]
+        test = treatment.strip() in self.instruments_data.columns
+        test = test & (treatment.strip() in self.data.columns)
+        if not test:
+            raise DataException(
+                f"""
+                The treatment variable:
+                {treatment} must appear in the instrument_data to be used
+                as an outcome variable and in the data object to be used as a covariate.
+                """
+            )

causalpy/tests/test_input_validation.py

@@ -202,3 +202,21 @@ def test_rd_validation_treated_is_dummy():
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
             treatment_threshold=0.5,
         )
+
+
+def test_iv_treatment_var_is_present():
+    data = pd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 5]})
+    instruments_formula = "risk  ~ 1 + logmort0"
+    formula = "loggdp ~  1 + risk"
+    instruments_data = pd.DataFrame({"z": [1, 3, 4], "w": [2, 3, 4]})
+
+    with pytest.raises(DataException):
+        _ = cp.pymc_experiments.InstrumentalVariable(
+            instruments_data=instruments_data,
+            data=data,
+            instruments_formula=instruments_formula,
+            formula=formula,
+            model=cp.pymc_models.InstrumentalVariableRegression(
+                sample_kwargs=sample_kwargs
+            ),
+        )