move plotting into experiment classes, removing PlotComponent entirely

Author: drbenvincent

causalpy/experiments/base.py

@@ -15,6 +15,11 @@
 Base class for quasi experimental designs.
 """
 
+from abc import abstractmethod
+
+from causalpy.pymc_models import PyMCModel
+from causalpy.skl_models import ScikitLearnModel
+
 
 class BaseExperiment:
     """Base class for quasi experimental designs."""
@@ -33,3 +38,26 @@ def idata(self):
     def print_coefficients(self, round_to=None):
         """Ask the model to print its coefficients."""
         self.model.print_coefficients(self.labels, round_to)
+
+    def plot(self, *args, **kwargs) -> tuple:
+        """Plot the model.
+
+        Internally, this function dispatches to either `bayesian_plot` or `ols_plot`
+        depending on the model type.
+        """
+        if isinstance(self.model, PyMCModel):
+            return self.bayesian_plot(*args, **kwargs)
+        elif isinstance(self.model, ScikitLearnModel):
+            return self.ols_plot(*args, **kwargs)
+        else:
+            raise ValueError("Unsupported model type")
+
+    @abstractmethod
+    def bayesian_plot(self, *args, **kwargs):
+        """Abstract method for plotting the model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def ols_plot(self, *args, **kwargs):
+        """Abstract method for plotting the model."""
+        raise NotImplementedError

causalpy/experiments/diff_in_diff.py

@@ -15,21 +15,26 @@
 Difference in differences
 """
 
+import arviz as az
 import numpy as np
 import pandas as pd
+import seaborn as sns
 from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 
 from causalpy.custom_exceptions import (
     DataException,
     FormulaException,
 )
+from causalpy.plot_utils import plot_xY
 from causalpy.pymc_models import PyMCModel
 from causalpy.skl_models import ScikitLearnModel
-from causalpy.utils import _is_variable_dummy_coded, convert_to_string
+from causalpy.utils import _is_variable_dummy_coded, convert_to_string, round_num
 
 from .base import BaseExperiment
 
+LEGEND_FONT_SIZE = 12
+
 
 class DifferenceInDifferences(BaseExperiment):
     """A class to analyse data from Difference in Difference settings.
@@ -205,18 +210,6 @@ def input_validation(self):
                 coded. Consisting of 0's and 1's only."""
             )
 
-    def plot(self, round_to=None) -> tuple[plt.Figure, plt.Axes]:
-        """
-        Plot the results
-
-        :param round_to:
-            Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers.
-        """
-        # Get a BayesianPlotComponent or OLSPlotComponent depending on the model
-        plot_component = self.model.get_plot_component()
-        fig, ax = plot_component.plot_difference_in_differences(self, round_to=round_to)
-        return fig, ax
-
     def summary(self, round_to=None) -> None:
         """Print summary of main results and model coefficients.
 
@@ -232,3 +225,216 @@ def summary(self, round_to=None) -> None:
     def _causal_impact_summary_stat(self, round_to=None) -> str:
         """Computes the mean and 94% credible interval bounds for the causal impact."""
         return f"Causal impact = {convert_to_string(self.causal_impact, round_to=round_to)}"
+
+    def bayesian_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
+        """
+        Plot the results
+
+        :param round_to:
+            Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers.
+        """
+        round_to = kwargs.get("round_to")
+
+        def _plot_causal_impact_arrow(results, ax):
+            """
+            draw a vertical arrow between `y_pred_counterfactual` and
+            `y_pred_counterfactual`
+            """
+            # Calculate y values to plot the arrow between
+            y_pred_treatment = (
+                results.y_pred_treatment["posterior_predictive"]
+                .mu.isel({"obs_ind": 1})
+                .mean()
+                .data
+            )
+            y_pred_counterfactual = (
+                results.y_pred_counterfactual["posterior_predictive"].mu.mean().data
+            )
+            # Calculate the x position to plot at
+            # Note that we force to be float to avoid a type error using np.ptp with boolean
+            # values
+            diff = np.ptp(
+                np.array(
+                    results.x_pred_treatment[results.time_variable_name].values
+                ).astype(float)
+            )
+            x = (
+                np.max(results.x_pred_treatment[results.time_variable_name].values)
+                + 0.1 * diff
+            )
+            # Plot the arrow
+            ax.annotate(
+                "",
+                xy=(x, y_pred_counterfactual),
+                xycoords="data",
+                xytext=(x, y_pred_treatment),
+                textcoords="data",
+                arrowprops={"arrowstyle": "<-", "color": "green", "lw": 3},
+            )
+            # Plot text annotation next to arrow
+            ax.annotate(
+                "causal\nimpact",
+                xy=(x, np.mean([y_pred_counterfactual, y_pred_treatment])),
+                xycoords="data",
+                xytext=(5, 0),
+                textcoords="offset points",
+                color="green",
+                va="center",
+            )
+
+        fig, ax = plt.subplots()
+
+        # Plot raw data
+        sns.scatterplot(
+            self.data,
+            x=self.time_variable_name,
+            y=self.outcome_variable_name,
+            hue=self.group_variable_name,
+            alpha=1,
+            legend=False,
+            markers=True,
+            ax=ax,
+        )
+
+        # Plot model fit to control group
+        time_points = self.x_pred_control[self.time_variable_name].values
+        h_line, h_patch = plot_xY(
+            time_points,
+            self.y_pred_control.posterior_predictive.mu,
+            ax=ax,
+            plot_hdi_kwargs={"color": "C0"},
+            label="Control group",
+        )
+        handles = [(h_line, h_patch)]
+        labels = ["Control group"]
+
+        # Plot model fit to treatment group
+        time_points = self.x_pred_control[self.time_variable_name].values
+        h_line, h_patch = plot_xY(
+            time_points,
+            self.y_pred_treatment.posterior_predictive.mu,
+            ax=ax,
+            plot_hdi_kwargs={"color": "C1"},
+            label="Treatment group",
+        )
+        handles.append((h_line, h_patch))
+        labels.append("Treatment group")
+
+        # Plot counterfactual - post-test for treatment group IF no treatment
+        # had occurred.
+        time_points = self.x_pred_counterfactual[self.time_variable_name].values
+        if len(time_points) == 1:
+            parts = ax.violinplot(
+                az.extract(
+                    self.y_pred_counterfactual,
+                    group="posterior_predictive",
+                    var_names="mu",
+                ).values.T,
+                positions=self.x_pred_counterfactual[self.time_variable_name].values,
+                showmeans=False,
+                showmedians=False,
+                widths=0.2,
+            )
+            for pc in parts["bodies"]:
+                pc.set_facecolor("C0")
+                pc.set_edgecolor("None")
+                pc.set_alpha(0.5)
+        else:
+            h_line, h_patch = plot_xY(
+                time_points,
+                self.y_pred_counterfactual.posterior_predictive.mu,
+                ax=ax,
+                plot_hdi_kwargs={"color": "C2"},
+                label="Counterfactual",
+            )
+            handles.append((h_line, h_patch))
+            labels.append("Counterfactual")
+
+        # arrow to label the causal impact
+        _plot_causal_impact_arrow(self, ax)
+
+        # formatting
+        ax.set(
+            xticks=self.x_pred_treatment[self.time_variable_name].values,
+            title=self._causal_impact_summary_stat(round_to),
+        )
+        ax.legend(
+            handles=(h_tuple for h_tuple in handles),
+            labels=labels,
+            fontsize=LEGEND_FONT_SIZE,
+        )
+        return fig, ax
+
+    def ols_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
+        """Generate plot for difference-in-differences"""
+        round_to = kwargs.get("round_to")
+        fig, ax = plt.subplots()
+
+        # Plot raw data
+        sns.lineplot(
+            self.data,
+            x=self.time_variable_name,
+            y=self.outcome_variable_name,
+            hue="group",
+            units="unit",
+            estimator=None,
+            alpha=0.25,
+            ax=ax,
+        )
+        # Plot model fit to control group
+        ax.plot(
+            self.x_pred_control[self.time_variable_name],
+            self.y_pred_control,
+            "o",
+            c="C0",
+            markersize=10,
+            label="model fit (control group)",
+        )
+        # Plot model fit to treatment group
+        ax.plot(
+            self.x_pred_treatment[self.time_variable_name],
+            self.y_pred_treatment,
+            "o",
+            c="C1",
+            markersize=10,
+            label="model fit (treament group)",
+        )
+        # Plot counterfactual - post-test for treatment group IF no treatment
+        # had occurred.
+        ax.plot(
+            self.x_pred_counterfactual[self.time_variable_name],
+            self.y_pred_counterfactual,
+            "go",
+            markersize=10,
+            label="counterfactual",
+        )
+        # arrow to label the causal impact
+        ax.annotate(
+            "",
+            xy=(1.05, self.y_pred_counterfactual),
+            xycoords="data",
+            xytext=(1.05, self.y_pred_treatment[1]),
+            textcoords="data",
+            arrowprops={"arrowstyle": "<->", "color": "green", "lw": 3},
+        )
+        ax.annotate(
+            "causal\nimpact",
+            xy=(
+                1.05,
+                np.mean([self.y_pred_counterfactual[0], self.y_pred_treatment[1]]),
+            ),
+            xycoords="data",
+            xytext=(5, 0),
+            textcoords="offset points",
+            color="green",
+            va="center",
+        )
+        # formatting
+        ax.set(
+            xlim=[-0.05, 1.1],
+            xticks=[0, 1],
+            xticklabels=["pre", "post"],
+            title=f"Causal impact = {round_num(self.causal_impact, round_to)}",
+        )
+        ax.legend(fontsize=LEGEND_FONT_SIZE)
+        return fig, ax