[ENH] Add run_predictive utility for probabilistic modeling

Introduces `run_predictive` in the API for running predictive sampling. Refactored existing code in `test_prob_factory` to use the new utility, enhancing modularity and reducing redundancy.

Author: Leguark

gempy_probability/__init__.py

@@ -1,4 +1,5 @@
 from .modules.model_definition.prob_model_factory import make_gempy_pyro_model, GemPyPyroModel
 from .modules import likelihoods
+from .api.model_runner import run_predictive
 
 from ._version import __version__

gempy_probability/api/model_runner/__init__.py

@@ -0,0 +1 @@
+from ._pyro_runner import run_predictive

gempy_probability/api/model_runner/_pyro_runner.py

@@ -0,0 +1,37 @@
+import numpy as np
+import os
+import pyro
+import pyro.distributions as dist
+import torch
+from arviz import InferenceData
+from pyro.distributions import Distribution
+
+import gempy as gp
+from gempy_engine.core.backend_tensor import BackendTensor
+import gempy_probability as gpp
+from gempy_engine.core.data.interpolation_input import InterpolationInput
+from pyro.infer import Predictive
+import pyro
+import arviz as az
+import matplotlib.pyplot as plt
+
+from pyro.infer import NUTS
+from pyro.infer import MCMC
+from pyro.infer.autoguide import init_to_mean
+
+
+def run_predictive(prob_model: gpp.GemPyPyroModel, geo_model: gp.data.GeoModel, 
+                   y_obs_list: torch.Tensor, n_samples: int, plot_trace:bool=False) -> az.InferenceData:
+    predictive = Predictive(
+        model=prob_model,
+        num_samples=n_samples
+    )
+    prior: dict[str, torch.Tensor] = predictive(geo_model, y_obs_list)
+    # print("Number of interpolations: ", geo_model.counter)
+
+    data: az.InferenceData = az.from_pyro(prior=prior)
+    if plot_trace:
+        az.plot_trace(data.prior)
+        plt.show()
+        
+    return data

gempy_probability/core/sampler.py

@@ -4,6 +4,7 @@
 from typing import Callable, Dict
 
 import gempy as gp
+from gempy_engine.core.backend_tensor import BackendTensor
 from gempy_probability.modules.forwards import run_gempy_forward
 
 GemPyPyroModel = Callable[[gp.data.GeoModel, torch.Tensor], None]
@@ -13,7 +14,7 @@ def make_gempy_pyro_model(
         *,
         priors: Dict[str, Distribution],
         set_interp_input_fn: Callable[
-            [Dict[str, torch.Tensor], gp.data.GeoModel],
+            [Dict[str, Distribution], gp.data.GeoModel],
             gp.data.InterpolationInput
         ],
         likelihood_fn: Callable[[gp.data.Solutions], Distribution],
@@ -95,6 +96,9 @@ def make_gempy_pyro_model(
     >>> # Now this can be used with Predictive or MCMC directly:
     >>> #   Predictive(pyro_model, num_samples=100)(geo_model, obs_tensor)
     """
+
+    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
+    
     def model(geo_model: gp.data.GeoModel, obs_data: torch.Tensor):
         # 1) Sample from the user‐supplied priors
         samples: Dict[str, torch.Tensor] = {}

gempy_probability/modules/model_definition/prob_model_factory.py

@@ -33,7 +33,7 @@ def test_prob_model_factory() -> None:
     # TODO: Convert this into an options preset
     geo_model.interpolation_options.uni_degree = 0
     geo_model.interpolation_options.mesh_extraction = False
-    geo_model.interpolation_options.sigmoid_slope = 1100.
+    geo_model.interpolation_options.sigmoid_slope = 1100
 
     # region Minimal grid for the specific likelihood function
     x_loc = 6000
@@ -77,7 +77,6 @@ def modify_z_for_surface_point1(
         geo_model: gp.data.GeoModel,
 ) -> InterpolationInput:
     # TODO: We can make a factory for this type of functions
-    
     prior_key = r'$\mu_{top}$'
 
     from gempy.modules.data_manipulation import interpolation_input_from_structural_frame
@@ -104,16 +103,13 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     from pyro.infer.autoguide import init_to_mean
 
     # region prior sampling
-    predictive = Predictive(
-        model=prob_model,
-        num_samples=50
+    prior = gpp.run_predictive(
+        prob_model=prob_model,
+        geo_model=geo_model,
+        y_obs_list=y_obs_list,
+        n_samples=50,
+        plot_trace=True
     )
-    prior = predictive(geo_model, y_obs_list)
-    # print("Number of interpolations: ", geo_model.counter)
-
-    data = az.from_pyro(prior=prior)
-    az.plot_trace(data.prior)
-    plt.show()
 
     # endregion
 
@@ -184,3 +180,5 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     print(f"Thickness mean: {posterior_thickness_mean}")
     print("MCMC convergence diagnostics:")
     print(f"Divergences: {float(data.sample_stats.diverging.sum())}")
+
+