allow prior variables (#672)

Author: aloctavodia

preliz/internal/distribution_helper.py

@@ -145,6 +145,7 @@ def num_kurtosis(dist):
     "ZeroInflatedPoisson": {"psi": 0.8, "mu": 4.5},
     "Truncated": {"lower": -10, "upper": 10},
     "Censored": {"lower": -10, "upper": 10},
+    "Dirichlet": {"alpha": [1.0, 1.0, 1.0]},
 }
 
 

preliz/internal/plot_helper.py

@@ -4,7 +4,15 @@
 
 try:
     from IPython import get_ipython
-    from ipywidgets import Checkbox, FloatSlider, FloatText, IntSlider, IntText, ToggleButton
+    from ipywidgets import (
+        Checkbox,
+        FloatSlider,
+        FloatText,
+        IntSlider,
+        IntText,
+        RadioButtons,
+        ToggleButton,
+    )
 except ImportError:
     pass
 
@@ -323,7 +331,7 @@ def get_boxes(name, value, lower, upper):
     return text
 
 
-def get_textboxes(signature, model):
+def get_textboxes(signature, model, kind_plot):
     textboxes = {}
     for name, param in signature.parameters.items():
         if isinstance(param.default, int | float):
@@ -399,6 +407,13 @@ def get_textboxes(signature, model):
         tooltip="Resample",
     )
 
+    textboxes["__kind__"] = RadioButtons(
+        options=["hist", "kde", "ecdf"],
+        value=kind_plot,
+        description="Kind",
+        disabled=False,
+    )
+
     return textboxes
 
 
@@ -432,8 +447,9 @@ def looper(*args, **kwargs):
     return looper
 
 
-def plot_repr(results, kind_plot, references, iterations, ax):
+def plot_repr(results, kind_plot, references, iterations, stats_kwargs, ax):
     alpha = max(0.01, 1 - iterations * 0.009)
+    stats_kwargs.setdefault("alpha", alpha)
 
     if kind_plot == "hist":
         if results[0].dtype.kind == "i":
@@ -442,31 +458,29 @@ def plot_repr(results, kind_plot, references, iterations, ax):
                 ax.set_xticks(bins + 0.5)
         else:
             bins = "auto"
-        ax.hist(
-            results.T,
-            alpha=alpha,
-            density=True,
-            color=["0.5"] * iterations,
-            bins=bins,
-            histtype="step",
-        )
-        ax.hist(
-            np.concatenate(results),
-            density=True,
-            bins=bins,
-            color="k",
-            ls="--",
-            histtype="step",
-        )
+
+        stats_kwargs.setdefault("bins", bins)
+        stats_kwargs.setdefault("density", True)
+        stats_kwargs.setdefault("histtype", "step")
+        stats_kwargs.setdefault("alpha", alpha)
+        stats_kwargs.setdefault("color", ["0.5"] * iterations)
+
+        ax.hist(results.T, **stats_kwargs)
+        stats_kwargs.pop("color")
+        stats_kwargs.pop("ls", None)
+        ax.hist(np.concatenate(results), color="k", ls="--", **stats_kwargs)
     elif kind_plot == "kde":
+        stats_kwargs.setdefault("color", "0.5")
         for result in results:
-            ax.plot(*kde(result), "0.5", alpha=alpha)
+            ax.plot(*kde(result), **stats_kwargs)
         ax.plot(*kde(np.concatenate(results)), "k--")
+
     elif kind_plot == "ecdf":
+        stats_kwargs.setdefault("color", "0.5")
         ax.plot(
             np.sort(results, axis=1).T,
             np.linspace(0, 1, len(results[0]), endpoint=False),
-            color="0.5",
+            **stats_kwargs,
         )
         a = np.concatenate(results)
         ax.plot(np.sort(a), np.linspace(0, 1, len(a), endpoint=False), "k--")

preliz/ppls/agnostic.py

@@ -222,7 +222,9 @@ def match_preliz_dist(all_dist_str, source, engine):
     return matches
 
 
-def ppl_plot_decorator(func, iterations, kind_plot, references, plot_func, engine):
+def ppl_plot_decorator(
+    func, iterations, kind_plot, references, plot_func, engine, group, var_name, stats_kwargs
+):
     def looper(*args, **kwargs):
         kwargs.pop("__resample__")
         x_min = kwargs.pop("__x_min__")
@@ -234,6 +236,7 @@ def looper(*args, **kwargs):
         else:
             auto = False
 
+        var_to_plot = var_name
         y_min = kwargs.pop("__y_min__", None)
         y_max = kwargs.pop("__y_max__", None)
         set_ylim = kwargs.pop("__set_ylim__", False)
@@ -244,6 +247,8 @@ def looper(*args, **kwargs):
         else:
             auto_ylim = False
 
+        kind = kwargs.pop("__kind__", kind_plot)
+
         if engine == "preliz":
             results = []
             for _ in range(iterations):
@@ -255,28 +260,30 @@ def looper(*args, **kwargs):
         elif engine == "bambi":
             model = func(*args, **kwargs)
             model.build()
+            if var_to_plot is None:
+                var_to_plot = model.observed_RVs[0].name
+
             with disable_pymc_sampling_logs():
                 idata = model.prior_predictive(iterations)
-            results = (
-                idata["prior_predictive"]
-                .stack(sample=("chain", "draw"))[model.response_component.response.name]
-                .values.T
-            )
+            results = idata[group].stack(sample=("chain", "draw"))[var_to_plot].values.T
+            if group == "prior":
+                results = np.atleast_2d(results)
 
         elif engine == "pymc":
             with func(*args, **kwargs) as model:
-                obs_name = model.observed_RVs[0].name
+                if var_to_plot is None:
+                    var_to_plot = model.observed_RVs[0].name
                 with disable_pymc_sampling_logs():
                     idata = sample_prior_predictive(samples=iterations)
-                results = (
-                    idata["prior_predictive"].stack(sample=("chain", "draw"))[obs_name].values.T
-                )
+                results = idata[group].stack(sample=("chain", "draw"))[var_to_plot].values.T
+            if group == "prior":
+                results = np.atleast_2d(results)
 
         _, ax = plt.subplots()
         ax.set_xlim(x_min, x_max, auto=auto)
         ax.set_ylim(y_min, y_max, auto=auto_ylim)
         if plot_func is None:
-            plot_repr(results, kind_plot, references, iterations, ax)
+            plot_repr(results, kind, references, iterations, stats_kwargs, ax)
         else:
             plot_func(results, ax)
 

preliz/predictive/predictive_explorer.py

@@ -12,13 +12,23 @@
 
 
 def predictive_explorer(
-    fmodel, samples=50, kind_plot="ecdf", references=None, plot_func=None, engine="auto"
+    fmodel,
+    samples=50,
+    kind_plot="ecdf",
+    references=None,
+    plot_func=None,
+    engine="auto",
+    group="prior_predictive",
+    var_name=None,
+    stats_kwargs=None,
 ):
     """
     Explore how changing parameters in the prior affects the prior predictive distribution.
 
-    Use this function to interactively explore how a prior predictive distribution changes when the
-    priors are changed.
+    Use this function to interactively explore how a prior predictive distribution changes
+    when the priors are changed. It also allows you to visualize how one prior changes when
+    another prior is changed, this can be useful for prior that are not set independently, but
+    are dependent on each other.
 
     Parameters
     ----------
@@ -40,13 +50,28 @@ def predictive_explorer(
         Library used to define the fmodel. Either `preliz`, `pymc` or `bambi`. Default is `auto`.
         The function will automatically select the appropriate library to use based on the fmodel
         provided.
+    group : str, optional
+        Which group to use. Ignored if the model is defined in `preliz`.
+        Defaults to "prior_predictive". You can also pass "prior".
+    var_name: str, optional
+        The name of the variable to plot. Ignored if the model is defined in `preliz`.
+        If "group=prior_predictive" it defaults to the first variable in `observed_RVs`.
+        For "prior" it defaults to the last variable in `free_RVs`.
+    stats_kwargs : dict, optional
+        Additional keyword arguments to pass to the statistics function.
+        Defaults to an empty dictionary.
     """
+    if stats_kwargs is None:
+        stats_kwargs = {}
     source, signature, engine = inspect_source(fmodel)
     model = parse_function_for_pred_textboxes(source, signature, engine)
-    textboxes = get_textboxes(signature, model)
-    new_fmodel = ppl_plot_decorator(fmodel, samples, kind_plot, references, plot_func, engine)
+    textboxes = get_textboxes(signature, model, kind_plot)
+    new_fmodel = ppl_plot_decorator(
+        fmodel, samples, kind_plot, references, plot_func, engine, group, var_name, stats_kwargs
+    )
     out = interactive_output(new_fmodel, textboxes)
     default_names = [
+        "__kind__",
         "__set_xlim__",
         "__x_min__",
         "__x_max__",