treatment_time input validation + associated test

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -54,7 +54,7 @@ class TimeSeriesExperiment(ExperimentalDesign):
     def __init__(
         self,
         data: pd.DataFrame,
-        treatment_time: int,
+        treatment_time: int | float | pd.Timestamp,
         formula: str,
         model=None,
         **kwargs,
@@ -81,6 +81,12 @@ def __init__(
         self.post_X = np.asarray(new_x)
         self.post_y = np.asarray(new_y)
 
+        # Input validation
+        if isinstance(data.index, pd.DatetimeIndex):
+            assert isinstance(
+                treatment_time, pd.Timestamp
+            ), "If data.index is DatetimeIndex, treatment_time must be pd.Timestamp."
+
         # DEVIATION FROM SKL EXPERIMENT CODE =============================
         # fit the model to the observed (pre-intervention) data
         COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.pre_X.shape[0])}

causalpy/tests/test_integration_pymc_examples.py

@@ -217,6 +217,31 @@ def test_sc_brexit():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
+@pytest.mark.integration
+def test_sc_brexit_input_error():
+    """Test that an error is raised if the data index is datetime and the treatment time
+    is not pd.Timestamp."""
+    with pytest.raises(AssertionError):
+        df = cp.load_data("brexit")
+        df["Time"] = pd.to_datetime(df["Time"])
+        df.set_index("Time", inplace=True)
+        df = df.iloc[df.index > "2009", :]
+        treatment_time = "2016 June 24"  # NOTE This is not of type pd.Timestamp
+        df = df.drop(["Japan", "Italy", "US", "Spain"], axis=1)
+        target_country = "UK"
+        all_countries = df.columns
+        other_countries = all_countries.difference({target_country})
+        all_countries = list(all_countries)
+        other_countries = list(other_countries)
+        formula = target_country + " ~ " + "0 + " + " + ".join(other_countries)
+        _ = cp.pymc_experiments.SyntheticControl(
+            df,
+            treatment_time,
+            formula=formula,
+            model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
+        )
+
+
 @pytest.mark.integration
 def test_ancova():
     df = cp.load_data("anova1")