fix plotting of control units for synthetic control + add test coverage

Author: drbenvincent

causalpy/expt_prepostfit.py

@@ -182,3 +182,25 @@ class SyntheticControl(PrePostFit):
     """
 
     expt_type = "SyntheticControl"
+
+    def plot(self, round_to=None, plot_predictors: bool = False):
+        """
+        Plot the results
+
+        :param round_to:
+            Number of decimals used to round results. Defaults to 2. Use "None" to
+            return raw numbers.
+        :param plot_predictors:
+            Whether to plot the control units as well. Defaults to False.
+        """
+        # Get a BayesianPlotComponent or OLSPlotComponent depending on the model
+        plot_component = self.model.get_plot_component()
+        fig, ax = plot_component.plot_pre_post(self)
+        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)
+            ax[0].plot(
+                self.datapost.index, self.post_X, "-", c=[0.8, 0.8, 0.8], zorder=1
+            )
+
+        return fig, ax

causalpy/plotting.py

@@ -35,7 +35,7 @@ class BayesianPlotComponent(PlotComponent):
     """Plotting component for Bayesian models."""
 
     @staticmethod
-    def plot_pre_post(results, round_to=None):
+    def plot_pre_post(results, round_to=None, counterfactual_label=None):
         """Generate plot for pre-post experiment types, such as Interrupted Time Series
         and Synthetic Control."""
         datapre = results.datapre

causalpy/tests/test_integration_pymc_examples.py

@@ -440,13 +440,21 @@ def test_sc():
     assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
     result.summary()
+
     fig, ax = result.plot()
     assert isinstance(fig, plt.Figure)
     # For multi-panel plots, ax should be an array of axes
     assert isinstance(ax, np.ndarray) and all(
         isinstance(item, plt.Axes) for item in ax
     ), "ax must be a numpy.ndarray of plt.Axes"
 
+    fig, ax = result.plot(plot_predictors=True)
+    assert isinstance(fig, plt.Figure)
+    # For multi-panel plots, ax should be an array of axes
+    assert isinstance(ax, np.ndarray) and all(
+        isinstance(item, plt.Axes) for item in ax
+    ), "ax must be a numpy.ndarray of plt.Axes"
+
 
 @pytest.mark.integration
 def test_sc_brexit():
@@ -485,6 +493,7 @@ def test_sc_brexit():
     assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
     result.summary()
+
     fig, ax = result.plot()
     assert isinstance(fig, plt.Figure)
     # For multi-panel plots, ax should be an array of axes

causalpy/tests/test_integration_skl_examples.py

@@ -133,6 +133,14 @@ def test_sc():
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.SyntheticControl)
     result.summary()
+
+    fig, ax = result.plot()
+    assert isinstance(fig, plt.Figure)
+    # For multi-panel plots, ax should be an array of axes
+    assert isinstance(ax, np.ndarray) and all(
+        isinstance(item, plt.Axes) for item in ax
+    ), "ax must be a numpy.ndarray of plt.Axes"
+
     fig, ax = result.plot()
     assert isinstance(fig, plt.Figure)
     # For multi-panel plots, ax should be an array of axes