addressing Juan's comments

Author: drbenvincent

causalpy/plot_utils.py

@@ -1,11 +1,12 @@
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import xarray as xr
 from matplotlib.collections import PolyCollection
+from matplotlib.lines import Line2D
 
 
 def plot_xY(
@@ -14,9 +15,8 @@ def plot_xY(
     ax: plt.Axes,
     plot_hdi_kwargs: Optional[Dict[str, Any]] = None,
     hdi_prob: float = 0.94,
-    label: Optional[str] = "",
-    include_label: bool = True,
-):
+    label: Union[str, None] = None,
+) -> Tuple[Line2D, PolyCollection]:
     """Utility function to plot HDI intervals."""
 
     if plot_hdi_kwargs is None:
@@ -27,33 +27,23 @@ def plot_xY(
         Y.mean(dim=["chain", "draw"]),
         ls="-",
         **plot_hdi_kwargs,
-        label=f"{label}" if include_label else None,
+        label=f"{label}",
     )
     ax_hdi = az.plot_hdi(
         x,
         Y,
         hdi_prob=hdi_prob,
         fill_kwargs={
             "alpha": 0.25,
-            "label": " ",  # f"{hdi_prob*100}% HDI" if include_label else None,
+            "label": " ",
         },
         smooth=False,
         ax=ax,
         **plot_hdi_kwargs,
     )
-    # Return handle to patch.
-    # We get a list of the childen of the axis
-    # Filter for just the PolyCollection objects
-    # Take the last one
+    # Return handle to patch. We get a list of the childen of the axis. Filter for just
+    # the PolyCollection objects. Take the last one.
     h_patch = list(
         filter(lambda x: isinstance(x, PolyCollection), ax_hdi.get_children())
     )[-1]
-
-    # if include_label:
-    #     handles, labels = ax.get_legend_handles_labels()
-    #     ax.legend(
-    #         handles=[(h1, h2) for h1, h2 in zip(handles[::2], handles[1::2])],
-    #         # labels=[l1 + " + " + l2 for l1, l2 in zip(labels[::2], labels[1::2])],
-    #         labels=[l1 for l1 in labels[::2]],
-    #     )
-    return h_line, h_patch
+    return (h_line, h_patch)

causalpy/pymc_experiments.py

@@ -122,7 +122,6 @@ def plot(self):
             self.datapre.index,
             self.pre_pred["posterior_predictive"].mu,
             ax=ax[0],
-            include_label=False,
             plot_hdi_kwargs={"color": "C0"},
         )
         handles = [(h_line, h_patch)]
@@ -137,7 +136,6 @@ def plot(self):
             self.datapost.index,
             self.post_pred["posterior_predictive"].mu,
             ax=ax[0],
-            include_label=False,
             # label="Synthetic control",
             plot_hdi_kwargs={"color": "C1"},
         )
@@ -152,9 +150,8 @@ def plot(self):
                 self.post_pred, group="posterior_predictive", var_names="mu"
             ).mean("sample"),
             y2=np.squeeze(self.post_y),
-            color="C2",
+            color="C0",
             alpha=0.25,
-            # label="Causal impact",
         )
         handles.append(h)
         labels.append("Causal impact")
@@ -171,21 +168,19 @@ def plot(self):
             self.datapre.index,
             self.pre_impact,
             ax=ax[1],
-            include_label=False,
             plot_hdi_kwargs={"color": "C0"},
         )
         plot_xY(
             self.datapost.index,
             self.post_impact,
             ax=ax[1],
-            include_label=False,
             plot_hdi_kwargs={"color": "C1"},
         )
         ax[1].axhline(y=0, c="k")
         ax[1].fill_between(
             self.datapost.index,
             y1=self.post_impact.mean(["chain", "draw"]),
-            color="C2",
+            color="C0",
             alpha=0.25,
             label="Causal impact",
         )
@@ -197,7 +192,6 @@ def plot(self):
             self.datapost.index,
             self.post_impact_cumulative,
             ax=ax[2],
-            include_label=False,
             plot_hdi_kwargs={"color": "C1"},
         )
         ax[2].axhline(y=0, c="k")
@@ -209,7 +203,6 @@ def plot(self):
                 ls="-",
                 lw=3,
                 color="r",
-                # label="Treatment time",
             )
 
         ax[0].legend(
@@ -434,7 +427,7 @@ def plot(self):
                 widths=0.2,
             )
             for pc in parts["bodies"]:
-                pc.set_facecolor("C2")
+                pc.set_facecolor("C0")
                 pc.set_edgecolor("None")
                 pc.set_alpha(0.5)
         else: