Add plot_interactive to Dirichlet (#332)

* Add xy_lim to plot_dirichlet

* Added plot_interactive to Dirichlet

* Tests for plot_interactive

* Change x to y, and renamed dim to alpha

Author: rohanbabbar04

preliz/distributions/continuous_multivariate.py

@@ -8,13 +8,18 @@
 from copy import copy
 
 import numpy as np
+
+try:
+    from ipywidgets import interactive, widgets
+except ImportError:
+    pass
 from scipy import stats
 
 from .distributions_multivariate import Continuous
 from .continuous import Beta, Normal
 from ..internal.distribution_helper import all_not_none
 from ..internal.plot_helper_multivariate import plot_dirichlet, plot_mvnormal
-
+from ..internal.plot_helper import check_inside_notebook, get_slider
 
 eps = np.finfo(float).eps
 
@@ -63,7 +68,7 @@ def __init__(self, alpha=None):
         self._parametrization(alpha)
 
     def _parametrization(self, alpha=None):
-        self.param_names = "alpha"
+        self.param_names = ("alpha",)
         self.params_support = ((eps, np.inf),)
 
         self.alpha = alpha
@@ -205,6 +210,108 @@ def plot_ppf(
             self, "ppf", "marginals", pointinterval, interval, levels, None, figsize, ax
         )
 
+    def plot_interactive(
+        self,
+        kind="pdf",
+        xy_lim="both",
+        pointinterval=True,
+        interval="hdi",
+        levels=None,
+        figsize=None,
+    ):
+        """
+        Interactive exploration of parameters
+
+        Parameters
+        ----------
+        kind : str:
+            Type of plot. Available options are `pdf`, `cdf` and `ppf`.
+        xy_lim : str or tuple
+            Set the limits of the x-axis and/or y-axis.
+            Defaults to `"both"`, the limits of both axes are fixed for all subplots.
+            Use `"auto"` for automatic rescaling of x-axis and y-axis.
+            Or set them manually by passing a tuple of 4 elements,
+            the first two for x-axis, the last two for y-axis. The tuple can have `None`.
+        pointinterval : bool
+            Whether to include a plot of the quantiles. Defaults to False.
+            If `True` the default is to plot the median and two inter-quantiles ranges.
+        interval : str
+            Type of interval. Available options are the highest density interval `"hdi"` (default),
+            equal tailed interval `"eti"` or intervals defined by arbitrary `"quantiles"`.
+        levels : list
+            Mass of the intervals. For hdi or eti the number of elements should be 2 or 1.
+            For quantiles the number of elements should be 5, 3, 1 or 0
+            (in this last case nothing will be plotted).
+        figsize : tuple
+            Size of the figure
+        """
+
+        check_inside_notebook()
+
+        args = dict(zip(self.param_names, self.params))
+        self.__init__(**args)  # pylint: disable=unnecessary-dunder-call
+        if kind == "pdf":
+            w_checkbox_marginals = widgets.Checkbox(
+                value=True,
+                description="marginals",
+                disabled=False,
+                indent=False,
+            )
+            plot_widgets = {"marginals": w_checkbox_marginals}
+        else:
+            plot_widgets = {}
+        for index, dim in enumerate(self.params[0]):
+            plot_widgets[f"alpha-{index + 1}"] = get_slider(
+                f"alpha-{index + 1}", dim, *self.params_support[0]
+            )
+
+        def plot(**args):
+            if kind == "pdf":
+                marginals = args.pop("marginals")
+            params = {"alpha": np.asarray(list(args.values()), dtype=float)}
+            self.__init__(**params)  # pylint: disable=unnecessary-dunder-call
+            if kind == "pdf":
+                plot_dirichlet(
+                    self,
+                    "pdf",
+                    marginals,
+                    pointinterval,
+                    interval,
+                    levels,
+                    "full",
+                    figsize,
+                    None,
+                    xy_lim,
+                )
+            elif kind == "cdf":
+                plot_dirichlet(
+                    self,
+                    "cdf",
+                    "marginals",
+                    pointinterval,
+                    interval,
+                    levels,
+                    "full",
+                    figsize,
+                    None,
+                    xy_lim,
+                )
+            elif kind == "ppf":
+                plot_dirichlet(
+                    self,
+                    "cdf",
+                    "marginals",
+                    pointinterval,
+                    interval,
+                    levels,
+                    None,
+                    figsize,
+                    None,
+                    xy_lim,
+                )
+
+        return interactive(plot, **plot_widgets)
+
 
 class MvNormal(Continuous):
     r"""

preliz/internal/plot_helper_multivariate.py

@@ -95,6 +95,7 @@ def plot_dirichlet(
     support,
     figsize,
     axes,
+    xy_lim="auto",
 ):
     """Plot pdf, cdf or ppf of Dirichlet distribution."""
 
@@ -104,6 +105,10 @@ def plot_dirichlet(
     if figsize is None:
         figsize = (12, 4)
 
+    if isinstance(xy_lim, tuple):
+        xlim = xy_lim[:2]
+        ylim = xy_lim[2:]
+
     if marginals:
         a_0 = alpha.sum()
         cols, rows = get_cols_rows(dim)
@@ -116,8 +121,18 @@ def plot_dirichlet(
                 ax.remove()
 
         for a_i, ax in zip(alpha, axes):
+            marginal_dist = dist.marginal(a_i, a_0 - a_i)
+            if xy_lim == "both":
+                xlim = marginal_dist._finite_endpoints("full")
+                xvals = marginal_dist.xvals("restricted")
+                if representation == "pdf":
+                    max_pdf = np.max(marginal_dist.pdf(xvals))
+                    ylim = (-max_pdf * 0.075, max_pdf * 1.5)
+                elif representation == "ppf":
+                    max_ppf = marginal_dist.ppf(0.999)
+                    ylim = (-max_ppf * 0.075, max_ppf * 1.5)
             if representation == "pdf":
-                dist.marginal(a_i, a_0 - a_i).plot_pdf(
+                marginal_dist.plot_pdf(
                     pointinterval=pointinterval,
                     interval=interval,
                     levels=levels,
@@ -126,7 +141,7 @@ def plot_dirichlet(
                     ax=ax,
                 )
             elif representation == "cdf":
-                dist.marginal(a_i, a_0 - a_i).plot_cdf(
+                marginal_dist.plot_cdf(
                     pointinterval=pointinterval,
                     interval=interval,
                     levels=levels,
@@ -135,14 +150,17 @@ def plot_dirichlet(
                     ax=ax,
                 )
             elif representation == "ppf":
-                dist.marginal(a_i, a_0 - a_i).plot_ppf(
+                marginal_dist.plot_ppf(
                     pointinterval=pointinterval,
                     interval=interval,
                     levels=levels,
                     legend=False,
                     ax=ax,
                 )
-
+            if xy_lim != "auto" and representation != "ppf":
+                ax.set_xlim(*xlim)
+            if xy_lim != "auto" and representation != "cdf":
+                ax.set_ylim(*ylim)
         fig.text(0.5, 1, repr_to_matplotlib(dist), ha="center", va="center")
 
     else:

preliz/tests/test_plots.py

@@ -69,6 +69,26 @@ def test_dirichlet_plot(kwargs):
     a_dist.plot_ppf(**kwargs)
 
 
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {},
+        {"xy_lim": "auto"},
+        {"pointinterval": True, "xy_lim": "auto"},
+        {"pointinterval": True, "levels": [0.1, 0.9], "xy_lim": "both"},
+        {"pointinterval": True, "interval": "eti", "levels": [0.9], "xy_lim": (0.3, 0.9, 0.6, 1)},
+        {"pointinterval": True, "interval": "quantiles", "xy_lim": "both"},
+        {"pointinterval": True, "interval": "quantiles", "levels": [0.1, 0.5, 0.9]},
+        {"pointinterval": False, "figsize": (4, 4)},
+    ],
+)
+def test_plot_interactive_dirichlet(kwargs):
+    a_dist = pz.Dirichlet([2, 1, 2])
+    a_dist.plot_interactive(kind="pdf", **kwargs)
+    a_dist.plot_interactive(kind="cdf", **kwargs)
+    a_dist.plot_interactive(kind="ppf", **kwargs)
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [