#58 add legend to TimeSeriesExperiment plot

Author: drbenvincent

causalpy/plot_utils.py

@@ -1,17 +1,27 @@
 import arviz as az
 
 
-def plot_xY(x, Y, ax, plot_hdi_kwargs=dict(), hdi_prob: float = 0.94) -> None:
+def plot_xY(
+    x, Y, ax, plot_hdi_kwargs=dict(), hdi_prob: float = 0.94, include_label: bool = True
+) -> None:
     """Utility function to plot HDI intervals."""
 
     Y = Y.stack(samples=["chain", "draw"]).T
     az.plot_hdi(
         x,
         Y,
         hdi_prob=hdi_prob,
-        fill_kwargs={"alpha": 0.25, "label": f"{hdi_prob*100}% HDI"},
+        fill_kwargs={
+            "alpha": 0.25,
+            "label": f"{hdi_prob*100}% HDI" if include_label else None,
+        },
         smooth=False,
         ax=ax,
         **plot_hdi_kwargs,
     )
-    ax.plot(x, Y.mean(dim="samples"), color="k", label="Posterior mean")
+    ax.plot(
+        x,
+        Y.mean(dim="samples"),
+        color="k",
+        label="Posterior mean" if include_label else None,
+    )

causalpy/pymc_experiments.py

@@ -89,18 +89,21 @@ def plot(self):
         plot_xY(
             self.datapre.index, self.pre_pred["posterior_predictive"].y_hat, ax=ax[0]
         )
-        ax[0].plot(self.datapre.index, self.pre_y, "k.")
+        ax[0].plot(self.datapre.index, self.pre_y, "k.", label="Observations")
         # post intervention period
         plot_xY(
-            self.datapost.index, self.post_pred["posterior_predictive"].y_hat, ax=ax[0]
+            self.datapost.index,
+            self.post_pred["posterior_predictive"].y_hat,
+            ax=ax[0],
+            include_label=False,
         )
         ax[0].plot(self.datapost.index, self.post_y, "k.")
         ax[0].set(
             title=f"Pre-intervention Bayesian $R^2$: {self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
         )
 
         plot_xY(self.datapre.index, self.pre_impact, ax=ax[1])
-        plot_xY(self.datapost.index, self.post_impact, ax=ax[1])
+        plot_xY(self.datapost.index, self.post_impact, ax=ax[1], include_label=False)
         ax[1].axhline(y=0, c="k")
         ax[1].set(title="Causal Impact")
 
@@ -115,8 +118,11 @@ def plot(self):
                 ls="-",
                 lw=3,
                 color="r",
-                label="treatment time",
+                label="Treatment time",
             )
+
+        ax[0].legend(fontsize=LEGEND_FONT_SIZE)
+
         return (fig, ax)
 
 

causalpy/skl_experiments.py

@@ -117,7 +117,7 @@ def plot(self):
                 ls="-",
                 lw=3,
                 color="r",
-                label="treatment time",
+                label="Treatment time",
             )
 
         ax[0].legend(fontsize=LEGEND_FONT_SIZE)