Add support for engine="auto" and bambi models in predictive explorer

preliz/internal/parser.py

@@ -4,6 +4,8 @@
 from sys import modules
 
 import numpy as np
+import pymc as pm
+import bambi as bmb
 
 from preliz import distributions
 from .distribution_helper import init_vals
@@ -13,8 +15,18 @@ def inspect_source(fmodel):
     source = inspect.getsource(fmodel)
     signature = inspect.signature(fmodel)
     source = re.sub(r"#.*$|^#.*$", "", source, flags=re.MULTILINE)
-
-    return source, signature
+    default_params = {
+        name: (param.default if param.default is not inspect.Parameter.empty else np.nan)
+        for name, param in signature.parameters.items()
+    }
+    model = fmodel(**default_params)
+    if isinstance(model, pm.Model):
+        engine = "pymc"
+    elif isinstance(model, bmb.Model):
+        engine = "bambi"
+    else:
+        engine = "preliz"
+    return source, signature, engine
 
 
 def parse_function_for_pred_textboxes(source, signature, engine="preliz"):

preliz/internal/plot_helper.py

@@ -442,7 +442,7 @@ def looper(*args, **kwargs):
         model = func(*args, **kwargs)
         model.build()
         with disable_pymc_sampling_logs():
-            idata = model.prior_predictive(iterations, random_seed=iterations)
+            idata = model.prior_predictive(iterations)
         results = extract(idata, group="prior_predictive")[model.response_name].values.T
 
         _, ax = plt.subplots()

preliz/internal/predictive_helper.py

@@ -1,6 +1,5 @@
 from sys import modules
 
-import re
 import numpy as np
 
 
@@ -69,12 +68,3 @@ def select_prior_samples(selected, prior_samples, model):
     subsample = {rv: prior_samples[rv][selected] for rv in model.keys()}
 
     return subsample
-
-
-def get_engine(source):
-    source = re.sub(r"#.*$|^#.*$", "", source, flags=re.MULTILINE)
-    if re.search(r"\s*(?:bmb|bambi)\.Model\s*", source):
-        return "bambi"
-    if re.search(r"\s*(?:pm|pymc)\.Model\s*", source):
-        return "pymc"
-    return "preliz"

preliz/predictive/predictive_explorer.py

@@ -10,7 +10,6 @@
     pymc_plot_decorator,
     bambi_plot_decorator,
 )
-from preliz.internal.predictive_helper import get_engine
 
 
 def predictive_explorer(
@@ -44,8 +43,7 @@ def predictive_explorer(
         The function will automatically select the appropriate library to use based on the fmodel
         provided.
     """
-    source, signature = inspect_source(fmodel)
-    engine = get_engine(source) if engine == "auto" else engine
+    source, signature, engine = inspect_source(fmodel)
     model = parse_function_for_pred_textboxes(source, signature, engine)
     textboxes = get_textboxes(signature, model)
     if engine == "pymc":