make get_plot_data methods public, links to functions in docstrings, add arviz to intersphinx_mapping

Author: drbenvincent

causalpy/experiments/base.py

@@ -83,22 +83,22 @@ def _ols_plot(self, *args, **kwargs):
     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`
+        Internally, this function dispatches to either :func:`get_plot_data_bayesian` or :func:`get_plot_data_ols`
         depending on the model type.
         """
         if isinstance(self.model, PyMCModel):
-            return self._get_plot_data_bayesian(*args, **kwargs)
+            return self.get_plot_data_bayesian(*args, **kwargs)
         elif isinstance(self.model, RegressorMixin):
-            return self._get_plot_data_ols(*args, **kwargs)
+            return self.get_plot_data_ols(*args, **kwargs)
         else:
             raise ValueError("Unsupported model type")
 
     @abstractmethod
-    def _get_plot_data_bayesian(self, *args, **kwargs):
+    def get_plot_data_bayesian(self, *args, **kwargs):
         """Abstract method for recovering plot data."""
-        raise NotImplementedError("_get_plot_data_bayesian method not yet implemented")
+        raise NotImplementedError("get_plot_data_bayesian method not yet implemented")
 
     @abstractmethod
-    def _get_plot_data_ols(self, *args, **kwargs):
+    def get_plot_data_ols(self, *args, **kwargs):
         """Abstract method for recovering plot data."""
-        raise NotImplementedError("_get_plot_data_ols method not yet implemented")
+        raise NotImplementedError("get_plot_data_ols method not yet implemented")

causalpy/experiments/prepostfit.py

@@ -303,9 +303,12 @@ def _ols_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, List[plt.Axes]
 
         return (fig, ax)
 
-    def _get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame:
+    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.
+
+        :param hdi_prob:
+            Prob for which the highest density interval will be computed. The default value is defined as the default from the :func:`arviz.hdi` function.
         """
         if isinstance(self.model, PyMCModel):
             hdi_pct = int(round(hdi_prob * 100))
@@ -350,7 +353,7 @@ def _get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame:
         else:
             raise ValueError("Unsupported model type")
 
-    def _get_plot_data_ols(self) -> pd.DataFrame:
+    def get_plot_data_ols(self) -> pd.DataFrame:
         """
         Recover the data of a PrePostFit experiment along with the prediction and causal impact information.
         """

docs/source/conf.py

@@ -109,6 +109,7 @@
 
 # -- intersphinx config -------------------------------------------------------
 intersphinx_mapping = {
+    "arviz": ("https://python.arviz.org/en/stable/", None),
     "examples": ("https://www.pymc.io/projects/examples/en/latest/", None),
     "mpl": ("https://matplotlib.org/stable", None),
     "numpy": ("https://numpy.org/doc/stable/", None),