refactor ppe and related functions

preliz/internal/optimization.py

@@ -259,44 +259,35 @@ def optimize_pymc_model(
     fmodel,
     target,
     num_draws,
-    bounds,
+    opt_iterations,
     initial_guess,
-    prior,
-    preliz_model,
-    transformed_var_info,
     rng,
 ):
-    for idx in range(401):
+    prior_array = np.zeros((opt_iterations, len(initial_guess)))
+
+    for idx in range(opt_iterations + 1):
         # can we sample systematically from these and less random?
         # This should be more flexible and allow other targets than just
         # a PreliZ distribution
         if isinstance(target, list):
             obs = get_weighted_rvs(target, num_draws, rng)
         else:
-            obs = target.rvs(num_draws, random_state=rng)
+            obs = target.rvs(num_draws, idx)
         result = minimize(
             fmodel,
             initial_guess,
             tol=0.001,
-            method="SLSQP",
-            args=(obs, transformed_var_info, preliz_model),
-            bounds=bounds,
+            method="powell",
+            args=(obs),
         )
-        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
+        # To help minimize the effect of priors we don't save the first result
+        # and instead we use it as the initial guess for the next optimization
+        # Updating the initial guess also helps reduce computational times
         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():
-        prior[key] = np.array(value)
+            prior_array[idx - 1] = result.x
+        initial_guess = result.x
 
-    return prior
+    return prior_array
 
 
 def relative_error(dist, lower, upper, required_mass):

preliz/ppls/agnostic.py

@@ -17,7 +17,11 @@
 from preliz.internal.plot_helper import plot_repr
 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.pymc_io import (
+    extract_preliz_distributions,
+    retrieve_variable_info,
+    write_pymc_string,
+)
 from preliz.ppls.bambi_io import (
     get_pymc_model,
     write_bambi_string,
@@ -65,14 +69,15 @@ def posterior_to_prior(model, idata, new_families=None, engine="auto"):
     if engine == "bambi":
         model = get_pymc_model(model)
 
-    _, _, preliz_model, _, untransformed_var_info, *_ = get_model_information(model)
+    preliz_model = extract_preliz_distributions(model)
+    var_info, _ = retrieve_variable_info(model)
 
     new_priors = back_fitting_idata(idata, preliz_model, new_families)
 
     if engine == "bambi":
-        new_model = write_bambi_string(new_priors, untransformed_var_info)
+        new_model = write_bambi_string(new_priors, var_info)
     elif engine == "pymc":
-        new_model = write_pymc_string(new_priors, untransformed_var_info)
+        new_model = write_pymc_string(new_priors, var_info)
 
     return new_model
 

preliz/ppls/bambi_io.py

@@ -2,9 +2,12 @@
 
 import importlib
 import inspect
+from copy import copy
 import re
 from sys import modules
 
+import numpy as np
+
 
 def get_pymc_model(model):
     if not model.built:
@@ -19,17 +22,36 @@ def write_bambi_string(new_priors, var_info):
     So the user can copy and paste, ideally with none to minimal changes.
     """
     header = "{\n"
+    variables = []
+    names = list(new_priors.keys())
     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}),\n'
+        idxs = var_info[key][-1]
+        if idxs:
+            for i in idxs:
+                nkey = names[i]
+                cp_dist = copy(new_priors[nkey])
+                cp_dist._fit_moments(np.mean(value.mean()), np.mean(value.std()))
+
+                dist_name, dist_params = repr(cp_dist).split("(")
+                size = var_info[nkey][1]
+                if size > 1:
+                    variables[
+                        i
+                    ] = f'"{nkey}" : bmb.Prior("{dist_name}", {dist_params}, shape={size}),\n'
+                else:
+                    variables[i] = f'"{nkey}" : bmb.Prior("{dist_name}", {dist_params}),\n'
         else:
-            header += f'"{key}" : bmb.Prior("{dist_name}", {dist_params}),\n'
-
-    header = header.rstrip(", ") + "}"
-    return header
+            dist_name, dist_params = repr(value).split("(")
+            dist_params = dist_params.rstrip(")")
+            size = var_info[key][1]
+            if size > 1:
+                variables.append(
+                    f'"{key}" : bmb.Prior("{dist_name}", {dist_params}, shape={size}),\n'
+                )
+            else:
+                variables.append(f'"{key}" : bmb.Prior("{dist_name}", {dist_params}),\n')
+
+    return "".join([header] + variables)
 
 
 def from_bambi(fmodel, draws):