Get plot data for prepostfit experiments

causalpy/experiments/base.py

@@ -17,6 +17,7 @@
 
 from abc import abstractmethod
 
+import pandas as pd
 from sklearn.base import RegressorMixin
 
 from causalpy.pymc_models import PyMCModel
@@ -78,3 +79,26 @@
     def ols_plot(self, *args, **kwargs):
         """Abstract method for plotting the model."""
         raise NotImplementedError("ols_plot method not yet implemented")
+
+    def get_plot_data(self, *args, **kwargs) -> pd.DataFrame:
+        """Recover the data of a PrePostFit experiment along with the prediction and causal impact information.
+
+        Internally, this function dispatches to either `get_plot_data_bayesian` or `get_plot_data_ols`
+        depending on the model type.
+        """
+        if isinstance(self.model, PyMCModel):
+            return self.get_plot_data_bayesian(*args, **kwargs)
+        elif isinstance(self.model, RegressorMixin):
+            return self.get_plot_data_ols(*args, **kwargs)
+        else:
+            raise ValueError("Unsupported model type")
+
+    @abstractmethod
+    def get_plot_data_bayesian(self, *args, **kwargs):
+        """Abstract method for recovering plot data."""
+        raise NotImplementedError("get_plot_data_bayesian method not yet implemented")
+
+    @abstractmethod
+    def get_plot_data_ols(self, *args, **kwargs):
+        """Abstract method for recovering plot data."""
+        raise NotImplementedError("get_plot_data_ols method not yet implemented")

causalpy/experiments/prepostfit.py

@@ -25,7 +25,7 @@
 from sklearn.base import RegressorMixin
 
 from causalpy.custom_exceptions import BadIndexException
-from causalpy.plot_utils import plot_xY
+from causalpy.plot_utils import get_hdi_to_df, plot_xY
 from causalpy.pymc_models import PyMCModel
 from causalpy.utils import round_num
 
@@ -303,6 +303,60 @@
 
         return (fig, ax)
 
+    def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame:
+        """
+        Recover the data of a PrePostFit experiment along with the prediction and causal impact information.
+        """
+        if isinstance(self.model, PyMCModel):
+            pre_data = self.datapre.copy()
+            post_data = self.datapost.copy()
+
+            pre_data["prediction"] = (
+                az.extract(self.pre_pred, group="posterior_predictive", var_names="mu")
+                .mean("sample")
+                .values
+            )
+            post_data["prediction"] = (
+                az.extract(self.post_pred, group="posterior_predictive", var_names="mu")
+                .mean("sample")
+                .values
+            )
+            pre_data[["pred_hdi_lower", "pred_hdi_upper"]] = get_hdi_to_df(
+                self.pre_pred["posterior_predictive"].mu, hdi_prob=hdi_prob
+            ).set_index(pre_data.index)
+            post_data[["pred_hdi_lower", "pred_hdi_upper"]] = get_hdi_to_df(
+                self.post_pred["posterior_predictive"].mu, hdi_prob=hdi_prob
+            ).set_index(post_data.index)
+
+            pre_data["impact"] = self.pre_impact.mean(dim=["chain", "draw"]).values
+            post_data["impact"] = self.post_impact.mean(dim=["chain", "draw"]).values
+            pre_data[["impact_hdi_lower", "impact_hdi_upper"]] = get_hdi_to_df(
+                self.pre_impact, hdi_prob=hdi_prob
+            ).set_index(pre_data.index)
+            post_data[["impact_hdi_lower", "impact_hdi_upper"]] = get_hdi_to_df(
+                self.post_impact, hdi_prob=hdi_prob
+            ).set_index(post_data.index)
+
+            self.plot_data = pd.concat([pre_data, post_data])
+
+            return self.plot_data
+        else:
+            raise ValueError("Unsupported model type")
+
+    def get_plot_data_ols(self) -> pd.DataFrame:
+        """
+        Recover the data of a PrePostFit experiment along with the prediction and causal impact information.
+        """
+        pre_data = self.datapre.copy()
+        post_data = self.datapost.copy()
+        pre_data["prediction"] = self.pre_pred
+        post_data["prediction"] = self.post_pred
+        pre_data["impact"] = self.pre_impact
+        post_data["impact"] = self.post_impact
+        self.plot_data = pd.concat([pre_data, post_data])
+
+        return self.plot_data
+
 
 class InterruptedTimeSeries(PrePostFit):
     """

causalpy/plot_utils.py

@@ -79,3 +79,24 @@
         filter(lambda x: isinstance(x, PolyCollection), ax_hdi.get_children())
     )[-1]
     return (h_line, h_patch)
+
+
+def get_hdi_to_df(
+    x: xr.DataArray,
+    hdi_prob: float = 0.94,
+) -> pd.DataFrame:
+    """
+    Utility function to calculate and recover HDI intervals.
+
+    :param x:
+        Xarray data array
+    :param hdi_prob:
+        The size of the HDI, default is 0.94
+    """
+    hdi = (
+        az.hdi(x, hdi_prob=hdi_prob)
+        .to_dataframe()
+        .unstack(level="hdi")
+        .droplevel(0, axis=1)
+    )
+    return hdi