Merge pull request #200 from pymc-labs/its-class

Add `InterruptedTimeSeries` class

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -134,7 +134,7 @@ def _input_validation(self, data, treatment_time):
                 "If data.index is not DatetimeIndex, treatment_time must be pd.Timestamp."  # noqa: E501
             )
 
-    def plot(self):
+    def plot(self, counterfactual_label="Counterfactual", **kwargs):
         """Plot the results"""
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
@@ -161,7 +161,7 @@ def plot(self):
             plot_hdi_kwargs={"color": "C1"},
         )
         handles.append((h_line, h_patch))
-        labels.append("Synthetic control")
+        labels.append(counterfactual_label)
 
         ax[0].plot(self.datapost.index, self.post_y, "k.")
         # Shaded causal effect
@@ -243,14 +243,20 @@ def summary(self):
         self.print_coefficients()
 
 
+class InterruptedTimeSeries(PrePostFit):
+    """Interrupted time series analysis"""
+
+    expt_type = "Interrupted Time Series"
+
+
 class SyntheticControl(PrePostFit):
     """A wrapper around the PrePostFit class"""
 
     expt_type = "Synthetic Control"
 
-    def plot(self, plot_predictors=False):
+    def plot(self, plot_predictors=False, **kwargs):
         """Plot the results"""
-        fig, ax = super().plot()
+        fig, ax = super().plot(counterfactual_label="Synthetic control", **kwargs)
         if plot_predictors:
             # plot control units as well
             ax[0].plot(self.datapre.index, self.pre_X, "-", c=[0.8, 0.8, 0.8], zorder=1)

causalpy/skl_experiments.py

@@ -74,7 +74,7 @@ def __init__(
         # cumulative impact post
         self.post_impact_cumulative = np.cumsum(self.post_impact)
 
-    def plot(self):
+    def plot(self, counterfactual_label="Counterfactual", **kwargs):
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
         ax[0].plot(self.datapre.index, self.pre_y, "k.")
@@ -84,7 +84,7 @@ def plot(self):
         ax[0].plot(
             self.datapost.index,
             self.post_pred,
-            label="counterfactual",
+            label=counterfactual_label,
             ls=":",
             c="k",
         )
@@ -95,7 +95,7 @@ def plot(self):
             self.datapost.index,
             self.post_impact,
             "k.",
-            label="counterfactual",
+            label=counterfactual_label,
         )
         ax[1].axhline(y=0, c="k")
         ax[1].set(title="Causal Impact")
@@ -151,12 +151,18 @@ def plot_coeffs(self):
         )
 
 
+class InterruptedTimeSeries(PrePostFit):
+    """Interrupted time series analysis"""
+
+    expt_type = "Interrupted Time Series"
+
+
 class SyntheticControl(PrePostFit):
     """A wrapper around the PrePostFit class"""
 
-    def plot(self, plot_predictors=False):
+    def plot(self, plot_predictors=False, **kwargs):
         """Plot the results"""
-        fig, ax = super().plot()
+        fig, ax = super().plot(counterfactual_label="Synthetic control", **kwargs)
         if plot_predictors:
             # plot control units as well
             ax[0].plot(self.datapre.index, self.pre_X, "-", c=[0.8, 0.8, 0.8], zorder=1)

causalpy/tests/test_integration_pymc_examples.py

@@ -167,14 +167,14 @@ def test_its_covid():
         .set_index("date")
     )
     treatment_time = pd.to_datetime("2020-01-01")
-    result = cp.pymc_experiments.SyntheticControl(
+    result = cp.pymc_experiments.InterruptedTimeSeries(
         df,
         treatment_time,
         formula="standardize(deaths) ~ 0 + standardize(t) + C(month) + standardize(temp)",  # noqa E501
         model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     assert isinstance(df, pd.DataFrame)
-    assert isinstance(result, cp.pymc_experiments.SyntheticControl)
+    assert isinstance(result, cp.pymc_experiments.InterruptedTimeSeries)
     assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]