Merge pull request #63 from pymc-labs/bunch-of-changes

Misc bunch of changes

Author: drbenvincent

README.md

@@ -55,7 +55,7 @@ This is appropriate when you have multiple units, one of which is treated. You b
 |--|--|
 | ![](img/synthetic_control_skl.svg) | ![](img/synthetic_control_pymc.svg) |
 
-> The data (treated and untreated units), pre-treatment model fit, and counterfactual (i.e. the synthetic control) are plotted (top). The Frequentist analysis shows the causal impact as a blue shaded region, but this is not shown in the Bayesian analysis to avoid a cluttered chart. Instead, the Bayesian analysis shows shaded Bayesian credible regions of the model fit and counterfactual. Also shown is the causal impact (middle) and cumulative causal impact (bottom).
+> The data (treated and untreated units), pre-treatment model fit, and counterfactual (i.e. the synthetic control) are plotted (top). The causal impact is shown as a blue shaded region. The Bayesian analysis shows shaded Bayesian credible regions of the model fit and counterfactual. Also shown is the causal impact (middle) and cumulative causal impact (bottom).
 
 ### Interrupted time series
 This is appropriate when you have a single treated unit, and therefore a single time series, and do _not_ have a set of untreated units.
@@ -71,7 +71,7 @@ This is appropriate when you have a single treated unit, and therefore a single
 |--|--|
 | ![](img/interrupted_time_series_skl.svg) | ![](img/interrupted_time_series_pymc.svg) |
 
-> The data, model fits, and counterfactual are plotted (top panels). The Frequentist analysis shows the causal impact with the blue shaded region, but this is not shown in the Bayesian analysis to avoid a cluttered chart. Instead, the Bayesian analysis shows shaded Bayesian credible regions of the model fits. Also shown is the causal impact (middle) and cumulative causal impact (bottom).
+> The data (treated and untreated units), pre-treatment model fit, and counterfactual (i.e. the synthetic control) are plotted (top). The causal impact is shown as a blue shaded region. The Bayesian analysis shows shaded Bayesian credible regions of the model fit and counterfactual. Also shown is the causal impact (middle) and cumulative causal impact (bottom).
 
 ### Difference in Differences
 

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,23 +89,46 @@ 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")
 
         ax[2].set(title="Cumulative Causal Impact")
         plot_xY(self.datapost.index, self.post_impact_cumulative, ax=ax[2])
+        ax[2].axhline(y=0, c="k")
+
+        # Shaded causal effect
+        ax[0].fill_between(
+            self.datapost.index,
+            y1=az.extract(
+                self.post_pred, group="posterior_predictive", var_names="y_hat"
+            ).mean("sample"),
+            y2=np.squeeze(self.post_y),
+            color="C0",
+            alpha=0.25,
+            label="Causal impact",
+        )
+        ax[1].fill_between(
+            self.datapost.index,
+            y1=self.post_impact.mean(["chain", "draw"]),
+            color="C0",
+            alpha=0.25,
+            label="Causal impact",
+        )
 
         # Intervention line
         for i in [0, 1, 2]:
@@ -114,8 +137,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

@@ -99,14 +99,14 @@ def plot(self):
             y2=np.squeeze(self.post_y),
             color="C0",
             alpha=0.25,
-            label="causal impact",
+            label="Causal impact",
         )
         ax[1].fill_between(
             self.datapost.index,
             y1=np.squeeze(self.post_impact),
             color="C0",
             alpha=0.25,
-            label="causal impact",
+            label="Causal impact",
         )
 
         # Intervention line
@@ -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)

docs/index.rst

@@ -8,6 +8,17 @@ CausalPy - Causal inference in quasi-experimental settings
 
 A Python package focussing on causal inference in quasi-experimental settings. The package allows for sophisticated Bayesian model fitting methods to be used in addition to traditional OLS.
 
+Support
+-------
+
+This repository is supported by `PyMC Labs <https://www.pymc-labs.io>`_.
+
+.. image:: ../img/pymc-labs-log.png
+   :align: center
+   :target: https://www.pymc-labs.io
+   :scale: 50 %
+
+
 .. toctree::
    :maxdepth: 2
    :caption: Contents:
@@ -33,9 +44,7 @@ Documentation outline
    api_plot_utils
 
 
-Indices and tables
-==================
+Index
+=====
 
 * :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search`