diff --git a/preliz/internal/optimization.py b/preliz/internal/optimization.py index 2da36fbb..ca54c5a7 100644 --- a/preliz/internal/optimization.py +++ b/preliz/internal/optimization.py @@ -253,30 +253,41 @@ def interval_short(params): def optimize_pymc_model( - fmodel, target, draws, prior, initial_guess, bounds, var_info, p_model, rng + fmodel, + target, + num_draws, + bounds, + initial_guess, + prior, + preliz_model, + transformed_var_info, + rng, ): - for _ in range(400): + for idx in range(401): # can we sample systematically from these and less random? # This should be more flexible and allow other targets than just - # a preliz distribution + # a PreliZ distribution if isinstance(target, list): - obs = get_weighted_rvs(target, draws, rng) + obs = get_weighted_rvs(target, num_draws, rng) else: - obs = target.rvs(draws, random_state=rng) + obs = target.rvs(num_draws, random_state=rng) result = minimize( fmodel, initial_guess, tol=0.001, method="SLSQP", - args=(obs, var_info, p_model), + args=(obs, transformed_var_info, preliz_model), bounds=bounds, ) - optimal_params = result.x + # To help minimize the effect of priors + # We don't save the first result and insteas we use it as the initial guess + # for the next optimization + # Updating the initial guess also helps to provides more spread samples initial_guess = optimal_params - - for key, param in zip(prior.keys(), optimal_params): - prior[key].append(param) + if idx: + for key, param in zip(prior.keys(), optimal_params): + prior[key].append(param) # convert to numpy arrays for key, value in prior.items(): diff --git a/preliz/internal/predictive_helper.py b/preliz/internal/predictive_helper.py index 28de554d..b7b3a1c2 100644 --- a/preliz/internal/predictive_helper.py +++ b/preliz/internal/predictive_helper.py @@ -9,7 +9,7 @@ from .distribution_helper import get_distributions -def back_fitting(model, subset, new_families=True): +def back_fitting_ppa(model, subset, new_families=True): """ Use MLE to fit a subset of the prior samples to the marginal prior distributions """ diff --git a/preliz/ppls/agnostic.py b/preliz/ppls/agnostic.py index 4c0af48f..03f07460 100644 --- a/preliz/ppls/agnostic.py +++ b/preliz/ppls/agnostic.py @@ -7,7 +7,7 @@ from preliz.distributions import Gamma, Normal, HalfNormal from preliz.unidimensional.mle import mle from preliz.ppls.pymc_io import get_model_information, write_pymc_string -from preliz.ppls.bambi_io import get_bmb_model_information, write_bambi_string +from preliz.ppls.bambi_io import get_pymc_model, write_bambi_string _log = logging.getLogger("preliz") @@ -41,10 +41,23 @@ def posterior_to_prior(model, idata, alternative=None, engine="auto"): """ _log.info(""""This is an experimental method under development, use with caution.""") engine = get_engine(model) if engine == "auto" else engine + if engine == "bambi": - _, _, model_info, _, var_info2, *_ = get_bmb_model_information(model) + model = get_pymc_model(model) + + _, _, preliz_model, _, untransformed_var_info, *_ = get_model_information(model) + + new_priors = back_fitting_idata(idata, preliz_model, alternative) + + if engine == "bambi": + new_model = write_bambi_string(new_priors, untransformed_var_info) elif engine == "pymc": - _, _, model_info, _, var_info2, *_ = get_model_information(model) + new_model = write_pymc_string(new_priors, untransformed_var_info) + + return new_model + + +def back_fitting_idata(idata, model_info, alternative): new_priors = {} posterior = idata.posterior.stack(sample=("chain", "draw")) @@ -66,10 +79,4 @@ def posterior_to_prior(model, idata, alternative=None, engine="auto"): idx, _ = mle(dists, posterior[var].values, plot=False) new_priors[var] = dists[idx[0]] - - if engine == "bambi": - new_model = write_bambi_string(new_priors, var_info2) - elif engine == "pymc": - new_model = write_pymc_string(new_priors, var_info2) - - return new_model + return new_priors diff --git a/preliz/ppls/bambi_io.py b/preliz/ppls/bambi_io.py index d1e20c33..5208591b 100644 --- a/preliz/ppls/bambi_io.py +++ b/preliz/ppls/bambi_io.py @@ -1,11 +1,11 @@ -from preliz.ppls.pymc_io import get_model_information +"""Functions to communicate with Bambi.""" -def get_bmb_model_information(model): +def get_pymc_model(model): if not model.built: model.build() pymc_model = model.backend.model - return get_model_information(pymc_model) + return pymc_model def write_bambi_string(new_priors, var_info): @@ -13,15 +13,15 @@ def write_bambi_string(new_priors, var_info): Return a string with the new priors for the Bambi model. So the user can copy and paste, ideally with none to minimal changes. """ - header = "{" + header = "{\n" for key, value in new_priors.items(): dist_name, dist_params = repr(value).split("(") dist_params = dist_params.rstrip(")") size = var_info[key][1] if size > 1: - header += f'"{key}" : bmb.Prior("{dist_name}", {dist_params}, shape={size}), ' + header += f'"{key}" : bmb.Prior("{dist_name}", {dist_params}, shape={size}),\n' else: - header += f'"{key}" : bmb.Prior("{dist_name}", {dist_params}), ' + header += f'"{key}" : bmb.Prior("{dist_name}", {dist_params}),\n' header = header.rstrip(", ") + "}" return header diff --git a/preliz/ppls/pymc_io.py b/preliz/ppls/pymc_io.py index e75a8f2a..89f262eb 100644 --- a/preliz/ppls/pymc_io.py +++ b/preliz/ppls/pymc_io.py @@ -17,29 +17,29 @@ from preliz.internal.distribution_helper import get_distributions -def backfitting(prior, p_model, var_info2): +def back_fitting_pymc(prior, preliz_model, untransformed_var_info): """ Fit the samples from prior into user provided model's prior. from the perspective of ppe "prior" is actually an approximated posterior but from the users perspective is its prior. - We need to "backfitted" because we can not use arbitrary samples as priors. + We need to "backfit" because we can not use arbitrary samples as priors. We need probability distributions. """ new_priors = {} - for key, size_inf in var_info2.items(): + for key, size_inf in untransformed_var_info.items(): if not size_inf[2]: size = size_inf[1] if size > 1: params = [] for i in range(size): value = prior[f"{key}__{i}"] - dist = p_model[key] + dist = preliz_model[key] dist._fit_mle(value) params.append(dist.params) dist._parametrization(*[np.array(x) for x in zip(*params)]) else: value = prior[key] - dist = p_model[key] + dist = preliz_model[key] dist._fit_mle(value) new_priors[key] = dist @@ -81,7 +81,7 @@ def get_pymc_to_preliz(): return pymc_to_preliz -def get_guess(model, free_rvs): +def get_initial_guess(model, free_rvs): """ Get initial guess for optimization routine. """ @@ -104,17 +104,32 @@ def get_guess(model, free_rvs): def get_model_information(model): # pylint: disable=too-many-locals """ - Get information from the PyMC model. - - This needs some love. We even have a variable named var_info, - and another one var_info2! + Get information from a PyMC model. + + Parameters + ---------- + model : a PyMC model + + Returns + ------- + bounds : a list of tuples with the support of each marginal distribution in the model + prior : a dictionary with a key for each marginal distribution in the model and an empty + list as value. This will be filled with the samples from a backfitting procedure. + preliz_model : a dictionary with a key for each marginal distribution in the model and the + corresponding PreliZ distribution as value + transformed_var_info : a dictionary with a key for each transformed variable in the model + and a tuple with the shape, size and the indexes of the non-constant parents as value + untransformed_var_info : same as `transformed_var_info` but the keys are untransformed + variable names + num_draws : the number of observed samples + free_rvs : a list with the free random variables in the model """ bounds = [] prior = {} - p_model = {} - var_info = {} - var_info2 = {} + preliz_model = {} + transformed_var_info = {} + untransformed_var_info = {} free_rvs = [] pymc_to_preliz = get_pymc_to_preliz() rvs_to_values = model.rvs_to_values @@ -128,13 +143,13 @@ def get_model_information(model): # pylint: disable=too-many-locals r_v.owner.op.name if r_v.owner.op.name else str(r_v.owner.op).split("RV", 1)[0].lower() ) dist = copy(pymc_to_preliz[name]) - p_model[r_v.name] = dist + preliz_model[r_v.name] = dist if nc_parents: idxs = [free_rvs.index(var_) for var_ in nc_parents] # the keys are the name of the (transformed) variable - var_info[rvs_to_values[r_v].name] = (shape, size, idxs) + transformed_var_info[rvs_to_values[r_v].name] = (shape, size, idxs) # the keys are the name of the (untransformed) variable - var_info2[r_v.name] = (shape, size, idxs) + untransformed_var_info[r_v.name] = (shape, size, idxs) else: free_rvs.append(r_v) @@ -147,13 +162,21 @@ def get_model_information(model): # pylint: disable=too-many-locals prior[r_v.name] = [] # the keys are the name of the (transformed) variable - var_info[rvs_to_values[r_v].name] = (shape, size, nc_parents) + transformed_var_info[rvs_to_values[r_v].name] = (shape, size, nc_parents) # the keys are the name of the (untransformed) variable - var_info2[r_v.name] = (shape, size, nc_parents) - - draws = model.observed_RVs[0].eval().size - - return bounds, prior, p_model, var_info, var_info2, draws, free_rvs + untransformed_var_info[r_v.name] = (shape, size, nc_parents) + + num_draws = model.observed_RVs[0].eval().size + + return ( + bounds, + prior, + preliz_model, + transformed_var_info, + untransformed_var_info, + num_draws, + free_rvs, + ) def write_pymc_string(new_priors, var_info): diff --git a/preliz/predictive/ppa.py b/preliz/predictive/ppa.py index 2705faaf..dde8f2fd 100644 --- a/preliz/predictive/ppa.py +++ b/preliz/predictive/ppa.py @@ -19,7 +19,7 @@ plot_pp_mean, ) from ..internal.parser import get_prior_pp_samples, from_preliz, from_bambi -from ..internal.predictive_helper import back_fitting, select_prior_samples +from ..internal.predictive_helper import back_fitting_ppa, select_prior_samples from ..distributions import Normal from ..distributions.distributions import Distribution @@ -386,7 +386,7 @@ def on_return_prior(self): if len(selected) > 4: subsample = select_prior_samples(selected, self.prior_samples, self.model) - string, _ = back_fitting(self.model, subsample, new_families=False) + string, _ = back_fitting_ppa(self.model, subsample, new_families=False) self.fig.clf() plt.text(0.05, 0.5, string, fontsize=14) diff --git a/preliz/predictive/ppe.py b/preliz/predictive/ppe.py index a5b134ae..ed4d9ad0 100644 --- a/preliz/predictive/ppe.py +++ b/preliz/predictive/ppe.py @@ -1,68 +1,118 @@ """Projective predictive elicitation.""" -import logging +import warnings import numpy as np from preliz.internal.optimization import optimize_pymc_model +from preliz.ppls.bambi_io import get_pymc_model, write_bambi_string +from preliz.ppls.agnostic import back_fitting_idata +from preliz.internal.parser import get_engine from preliz.ppls.pymc_io import ( get_model_information, - get_guess, + get_initial_guess, compile_logp, - backfitting, + back_fitting_pymc, write_pymc_string, ) -_log = logging.getLogger("preliz") - - -def ppe(model, target, seed=0): +def ppe(model, target, method="projective", engine="auto", random_state=0): """ - Projective Predictive Elicitation. + Prior Predictive Elicitation. - This is an experimental method under development, use with caution. - Most likely thing will break. + This method is experimental and under development. It does not offers guarantees of + correctness. Use with caution and triple-check the results. Parameters ---------- model : a probabilistic model Currently it only works with PyMC model. More PPls coming soon. - target : a Preliz distribution or list + 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 from the user, - possibly using other Preliz's methods to obtain this distribution, such as maxent, + 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. - seed : int - A seed to initialize the Random Generator. Default is 0. + random_state : {None, int, numpy.random.Generator, numpy.random.RandomState} + Defaults to 0. Ignored if `method` is `"pathfinder"`. Returns ------- - prior : a dictionary - Prior samples approximating the prior distribution that will induce - a prior predictive distribution close to ``target``. - new_priors : a dictionary - PreliZ Distributions approximating the prior distribution inducing a - prior predictive distribution close to ``target``. - pymc_string : a string - This is the PyMC model string with the new priors. - Computed by taking the "prior samples" and fit it into the model's prior families - using MLE. + new_priors : str + A string representation of the new priors. The user can copy and paste it into + the model's code. Ideally, with none to minimal changes. """ - _log.info(""""This is an experimental method under development, use with caution.""") + warnings.warn( + """This method is experimental and under development with no guarantees of correctness. + Use with caution and triple-check the results.""" + ) + + rng = np.random.default_rng(random_state) + engine = get_engine(model) if engine == "auto" else engine + + # Get models information + if engine == "bambi": + model = get_pymc_model(model) + ( + bounds, + prior, + preliz_model, + transformed_var_info, + untransformed_var_info, + num_draws, + free_rvs, + ) = get_model_information(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, free_rvs) + # compile PyMC model + fmodel = compile_logp(model) + projection = optimize_pymc_model( + fmodel, + target, + num_draws, + bounds, + initial_guess, + prior, + preliz_model, + transformed_var_info, + rng, + ) + # Backfit `projected_posterior` into the model's prior-families + new_priors = back_fitting_pymc(projection, preliz_model, untransformed_var_info) + if engine == "bambi": + return write_bambi_string(new_priors, untransformed_var_info) + if engine == "pymc": + return write_pymc_string(new_priors, untransformed_var_info) + + # Fit the samples to the original prior distribution + # or to a set of predefined distributions + elif method == "pathfinder": + from pymc_experimental import fit # pylint:disable=import-outside-toplevel + + with model: + idata = fit(method="pathfinder", num_samples=1000) - # Get information from PyMC model - rng = np.random.default_rng(seed) - bounds, prior, p_model, var_info, var_info2, draws, free_rvs = get_model_information(model) - # Initial point for optimization - guess = get_guess(model, free_rvs) - # compile PyMC model - fmodel = compile_logp(model) - # find prior that induce a prior predictive distribution close to target - prior = optimize_pymc_model(fmodel, target, draws, prior, guess, bounds, var_info, p_model, rng) - # Fit the prior into the model's prior - # So we can write it as a PyMC model - new_priors = backfitting(prior, p_model, var_info2) + new_priors = back_fitting_idata(idata, preliz_model, alternative=False) + if engine == "bambi": + new_model = write_bambi_string(new_priors, untransformed_var_info) + elif engine == "pymc": + new_model = write_pymc_string(new_priors, untransformed_var_info) - return prior, new_priors, write_pymc_string(new_priors, var_info2) + return new_model diff --git a/preliz/tests/test_ppe.py b/preliz/tests/test_ppe.py index 39899aa3..2f53fc79 100644 --- a/preliz/tests/test_ppe.py +++ b/preliz/tests/test_ppe.py @@ -17,9 +17,8 @@ "sigma_z": 10, "target": pz.Normal(mu=174, sigma=20), "X": np.random.normal(0, 10, 120), - "new_mu_x": 173.998287, - "new_sigma_x": 1.946641, - "new_sigma_z": 19.926345, + "new_x": 174, + "new_z": 2.9, }, { "mu_x": [0, 1], @@ -27,9 +26,8 @@ "sigma_z": 10, "target": pz.Normal(mu=174, sigma=20), "X": np.random.normal(0, 10, 120), - "new_mu_x": (174.089111, 173.908702), - "new_sigma_x": (2.679288, 2.760424), - "new_sigma_z": 19.83069, + "new_x": [174, 174], + "new_z": 2.9, }, { "mu_x": 0, @@ -37,13 +35,12 @@ "sigma_z": 10, "target": [(pz.Normal(mu=174, sigma=20), 0.5), (pz.Normal(mu=176, sigma=19.5), 0.5)], "X": np.random.normal(0, 10, 120), - "new_mu_x": 174.852283, - "new_sigma_x": 1.794118, - "new_sigma_z": 19.683396, + "new_x": 175, + "new_z": 2.9, }, { "mu_x": [0, 1], - "sigma_x": [10, 10], + "sigma_x": 10, "sigma_z": 10, "target": [ (pz.Normal(mu=174, sigma=20), 0.5), @@ -51,25 +48,28 @@ (pz.StudentT(mu=174, sigma=20, nu=3), 0.1), ], "X": np.random.normal(0, 10, 120), - "new_mu_x": (174.75759, 174.775254), - "new_sigma_x": (2.728134, 2.978235), - "new_sigma_z": 21.247561, + "new_x": [175, 175], + "new_z": 3, }, ], ) def test_ppe(params): - Y = 2 + np.random.normal(params["X"], 1) + Y = np.zeros_like(params["X"]) target = params.pop("target") with pm.Model() as model: - x = pm.Normal("x", mu=params["mu_x"], sigma=params["sigma_x"]) - z = pm.HalfNormal("z", params["sigma_z"]) + x = pm.Normal("x", shape=1 if isinstance(params["mu_x"], int) else 2) + z = pm.HalfNormal("z") x_idx = ( - x - if max(np.asarray(params["mu_x"]).size, np.asarray(params["sigma_x"]).size) == 1 - else x[np.repeat(np.arange(2), Y.size / 2)] + x if np.asarray(params["mu_x"]).size == 1 else x[np.repeat(np.arange(2), Y.size / 2)] ) pm.Normal("y", x_idx, z, observed=Y) - _, new_prior, _ = pz.ppe(model, target) - assert_allclose(new_prior["x"].mu, params["new_mu_x"], 1) - assert_allclose(new_prior["x"].sigma, params["new_sigma_x"], 1) - assert_allclose(new_prior["z"].sigma, params["new_sigma_z"], 1) + + new_prior = ( + pz.ppe(model, target, method="projective").replace("\x1b[1m", "").replace("\x1b[0m", "") + ) + exec_context = {} + exec(new_prior, globals(), exec_context) + model = exec_context.get("model") + initial = model.initial_point() + assert_allclose(initial["x"], params["new_x"]) + assert_allclose(initial["z_log__"], params["new_z"], atol=0.1)