PPE: add alternative family option, remove "method" (#673)

* add families, clean code

* fix test

Author: aloctavodia

preliz/ppls/pymc_io.py

@@ -14,10 +14,12 @@
 except ModuleNotFoundError:
     pass
 
+from preliz.distributions import Gamma, HalfNormal, Normal
 from preliz.internal.distribution_helper import get_distributions
+from preliz.unidimensional import mle
 
 
-def back_fitting_pymc(prior, preliz_model, var_info):
+def back_fitting_pymc(prior, preliz_model, var_info, new_families=None):
     """
     Fit the samples from prior into user provided model's prior.
 
@@ -32,20 +34,36 @@ def back_fitting_pymc(prior, preliz_model, var_info):
             params = []
             for i in range(size):
                 opt_values = prior[rv_name][:, i]
+                # Not sure how to fit alternative families.
                 dist = preliz_model[rv_name]
                 dist._fit_mle(opt_values)
                 params.append(dist.params)
             dist._parametrization(*[np.array(x) for x in zip(*params)])
         else:
             opt_values = prior[rv_name]
-            dist = preliz_model[rv_name]
-            dist._fit_mle(opt_values)
+            dists = set_families(preliz_model[rv_name], rv_name, new_families)
+            idx, _ = mle(dists, opt_values, plot=False)
+            dist = dists[idx[0]]
 
         new_priors[rv_name] = dist
 
     return new_priors
 
 
+def set_families(dist, var, new_families):
+    dists = [dist]
+    if new_families is not None:
+        if new_families == "auto":
+            alt = [Normal(), HalfNormal(), Gamma()]
+            dists += [a for a in alt if dist.__class__.__name__ != a.__class__.__name__]
+        elif isinstance(new_families, list):
+            dists += new_families
+        elif isinstance(new_families, dict):
+            dists += new_families.get(var, [])
+
+    return dists
+
+
 def compile_mllk(model):
     """
     Compile the log-likelihood for a pymc model.

preliz/predictive/ppe.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from preliz.internal.optimization import optimize_pymc_model
-from preliz.ppls.agnostic import back_fitting_idata, get_engine
+from preliz.ppls.agnostic import get_engine
 from preliz.ppls.bambi_io import get_pymc_model, write_bambi_string
 from preliz.ppls.pymc_io import (
     back_fitting_pymc,
@@ -18,36 +18,39 @@
 )
 
 
-def ppe(model, target, method="projective", engine="auto", random_state=0):
+def ppe(model, target, engine="auto", new_families=None, random_state=0):
     """
     Prior Predictive Elicitation.
 
     This method is experimental and under development. It does not offers guarantees of
     correctness. Use with caution and triple-check the results.
 
+    With the projective method we attempt to find a prior that induces
+    a prior predictive distribution as close as possible to the target distribution
+
     Parameters
     ----------
     model : a probabilistic model
         Currently it only works with PyMC model. More PPls coming soon.
-    method : str
-        Method used to generate samples that match the target distribution.
-        Defaults to `"projective"`, another option is `"pathfinder"`.
-        If `"projective"`, the parameters of the priors are only used to provide an initial
-        guess for the optimization routine. Thus their effect on the result is smaller than in
-        traditional Bayesian inference, unless the priors are very vague or very strong.
-        If `"projective"`, the observed values are ignored, but not their size.
-        Pathfinder is a variational inference method so the role of the priors and observed values
-        is what is expected in Bayesian inference.
-    engine : str
-        Library used to define the model. Either `"auto"` (default), `"pymc"` or `"bambi"`.
-        Ig `"auto"`, the library is automatically detected.
     target : a PreliZ distribution or list
         Instance of a PreliZ distribution or a list of tuples where each tuple contains a PreliZ
         distribution and a weight.
         This represents the prior predictive distribution **previously** elicited by the user,
         possibly using other PreliZ's methods to obtain this distribution, such as maxent,
         roulette, quartile, etc.
         This should represent the domain-knowledge of the user and not any observed dataset.
+    engine : str
+        Library used to define the model. Either `"auto"` (default), `"pymc"` or `"bambi"`.
+        Ig `"auto"`, the library is automatically detected.
+    new_families : "auto", list or dict
+        Defaults to None, the samples are fit to the original prior distribution.
+        If "auto", the method evaluates the fit to the original prior plus a set of
+        predefined distributions.
+        Use a list of PreliZ distribution to specify the alternative distributions
+        you want to consider.
+        Use a dict with variables names in ``model`` as keys and a list of PreliZ
+        distributions as values. This allows to specify alternative distributions
+        per variable.
     random_state : {None, int, numpy.random.Generator, numpy.random.RandomState}
         Defaults to 0. Ignored if `method` is `"pathfinder"`.
 
@@ -73,44 +76,31 @@ def ppe(model, target, method="projective", engine="auto", random_state=0):
     preliz_model = extract_preliz_distributions(model)
     var_info, num_draws = retrieve_variable_info(model)
 
-    # With the projective method we attempt to find a prior that induces
-    # a prior predictive distribution as close as possible to the target distribution
-    if method == "projective":
-        # Initial point for optimization
-        initial_guess = get_initial_guess(model)
-        # compile PyMC model
-        fmodel, old_y_value, obs_rvs = compile_mllk(model)
-        projection_raveled = optimize_pymc_model(
-            fmodel,
-            target,
-            num_draws,
-            opt_iterations,
-            initial_guess,
-            rng,
-        )
-        # restore obs_rvs value in the model
-        model.rvs_to_values[obs_rvs] = old_y_value
-
-        projection_unraveled = unravel_projection(projection_raveled, var_info, opt_iterations)
-
-        # Backfit `projected_posterior` into the model's prior-families
-        projection_backfitted = back_fitting_pymc(projection_unraveled, preliz_model, var_info)
-
-        if engine == "bambi":
-            new_priors = write_bambi_string(projection_backfitted, var_info)
-        elif engine == "pymc":
-            new_priors = write_pymc_string(projection_backfitted, var_info)
-
-    elif method == "pathfinder":
-        from pymc_experimental import fit
-
-        with model:
-            idata = fit(method="pathfinder", num_samples=opt_iterations)
-
-        projection_backfitted = back_fitting_idata(idata, preliz_model, new_families=False)
-        if engine == "bambi":
-            new_priors = write_bambi_string(projection_backfitted, var_info)
-        elif engine == "pymc":
-            new_priors = write_pymc_string(projection_backfitted, var_info)
+    # Initial point for optimization
+    initial_guess = get_initial_guess(model)
+    # compile PyMC model
+    fmodel, old_y_value, obs_rvs = compile_mllk(model)
+    projection_raveled = optimize_pymc_model(
+        fmodel,
+        target,
+        num_draws,
+        opt_iterations,
+        initial_guess,
+        rng,
+    )
+    # restore obs_rvs value in the model
+    model.rvs_to_values[obs_rvs] = old_y_value
+
+    projection_unraveled = unravel_projection(projection_raveled, var_info, opt_iterations)
+
+    # Backfit `projected_posterior` into the model's prior-families
+    projection_backfitted = back_fitting_pymc(
+        projection_unraveled, preliz_model, var_info, new_families
+    )
+
+    if engine == "bambi":
+        new_priors = write_bambi_string(projection_backfitted, var_info)
+    elif engine == "pymc":
+        new_priors = write_pymc_string(projection_backfitted, var_info)
 
     return new_priors

preliz/tests/test_ppe.py

@@ -68,9 +68,7 @@ def test_ppe(params):
         )
         pm.Normal("y", x_idx, z, observed=Y)
 
-    new_prior = (
-        pz.ppe(model, target, method="projective").replace("\x1b[1m", "").replace("\x1b[0m", "")
-    )
+    new_prior = pz.ppe(model, target).replace("\x1b[1m", "").replace("\x1b[0m", "")
     exec_context = {}
     exec(new_prior, globals(), exec_context)
     model = exec_context.get("model")