[WIP/CLN/ENH] Extracting samplers to the api to make it easier

Leguark · Leguark · commit 0a151ede7ffc · 2025-05-27T16:38:59.000+01:00
diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py
@@ -1,5 +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 .api.model_runner import run_predictive, run_mcmc_for_NUTS, run_nuts_inference
 
 from ._version import __version__
diff --git a/gempy_probability/api/model_runner/__init__.py b/gempy_probability/api/model_runner/__init__.py
@@ -1 +1 @@
-from ._pyro_runner import run_predictive
+from ._pyro_runner import run_predictive, run_nuts_inference, run_mcmc_for_NUTS
diff --git a/gempy_probability/api/model_runner/_pyro_runner.py b/gempy_probability/api/model_runner/_pyro_runner.py
@@ -19,6 +19,8 @@
 from pyro.infer import MCMC
 from pyro.infer.autoguide import init_to_mean
 
+from gempy_probability.core.samplers_data import NUTSConfig
+
 
 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:
@@ -35,3 +37,65 @@ def run_predictive(prob_model: gpp.GemPyPyroModel, geo_model: gp.data.GeoModel,
         plt.show()
         
     return data
+
+def run_nuts_inference(prob_model: gpp.GemPyPyroModel, geo_model: gp.data.GeoModel, 
+                       y_obs_list: torch.Tensor, config: NUTSConfig, plot_trace:bool=False,
+                        run_posterior_predictive:bool=False) -> az.InferenceData:
+    
+    nuts_kernel = NUTS(
+        prob_model,
+        step_size=config.step_size,
+        adapt_step_size=config.adapt_step_size,
+        target_accept_prob=config.target_accept_prob,
+        max_tree_depth=config.max_tree_depth,
+        init_strategy=init_to_mean
+    )
+    data = run_mcmc_for_NUTS(
+        geo_model=geo_model,
+        nuts_kernel=nuts_kernel,
+        prob_model=prob_model,
+        y_obs_list=y_obs_list,
+        num_samples=config.num_samples,
+        warmup_steps=config.warmup_steps,
+        plot_trace=plot_trace,
+        run_posterior_predictive=run_posterior_predictive,
+    )
+
+    return data
+
+
+def run_mcmc_for_NUTS(
+        geo_model: gp.data.GeoModel,
+        nuts_kernel: NUTS,
+        prob_model: gpp.GemPyPyroModel,
+        y_obs_list: torch.Tensor,
+        num_samples: int,
+        warmup_steps: int,
+        plot_trace: bool = False,
+        run_posterior_predictive: bool = False,
+) -> az.InferenceData:
+    mcmc = MCMC(
+        kernel=nuts_kernel,
+        num_samples=num_samples,
+        warmup_steps=warmup_steps,
+        disable_validation=False,
+    )
+    mcmc.run(geo_model, y_obs_list)
+
+    if run_posterior_predictive:
+        posterior_predictive = Predictive(
+            model=prob_model,
+            posterior_samples=mcmc.get_samples(),
+        )(geo_model, y_obs_list)
+        data: az.InferenceData = az.from_pyro(
+            posterior=mcmc,
+            posterior_predictive=posterior_predictive,
+        )
+    else:
+        data: az.InferenceData = az.from_pyro(posterior=mcmc)
+
+    if plot_trace:
+        az.plot_trace(data)
+        plt.show()
+
+    return data
diff --git a/gempy_probability/core/samplers_data.py b/gempy_probability/core/samplers_data.py
@@ -2,7 +2,7 @@
 
 
 @dataclasses.dataclass
-class NUTS_Args:
+class NUTSConfig:
     step_size: float = 0.0085
     adapt_step_size: bool = True
     target_accept_prob: float = 0.9
diff --git a/gempy_probability/modules/forwards/_gempy_fw.py b/gempy_probability/modules/forwards/_gempy_fw.py
@@ -5,7 +5,7 @@
 
 
 def run_gempy_forward(interp_input: InterpolationInput, geo_model: gp.data.GeoModel) -> gp.data.Solutions:
-    BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
+    
     # compute model
     sol = gempy_engine.compute_model(
         interpolation_input=interp_input,
diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py
@@ -9,11 +9,13 @@
 from gempy_engine.core.backend_tensor import BackendTensor
 import gempy_probability as gpp
 from gempy_engine.core.data.interpolation_input import InterpolationInput
+from gempy_probability.core.samplers_data import NUTSConfig
 
 seed = 123456
 torch.manual_seed(seed)
 pyro.set_rng_seed(seed)
 
+
 def test_prob_model_factory() -> None:
     current_dir = os.path.dirname(os.path.abspath(__file__))
     data_path = os.path.abspath(os.path.join(current_dir, '..', '..', 'examples', 'tutorials', 'data'))
@@ -71,19 +73,20 @@ def test_prob_model_factory() -> None:
         prob_model=pyro_gravity_model,
         y_obs_list=torch.tensor([200, 210, 190])
     )
-    
+
+
 def modify_z_for_surface_point1(
         samples: dict[str, Distribution],
         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
     interp_input = interpolation_input_from_structural_frame(geo_model)
     new_tensor: torch.Tensor = torch.index_put(
         input=interp_input.surface_points.sp_coords,
-        indices=(torch.tensor([0]), torch.tensor([2])), # * This has to be Tensors
+        indices=(torch.tensor([0]), torch.tensor([2])),  # * This has to be Tensors
         values=(samples[prior_key])
     )
     interp_input.surface_points.sp_coords = new_tensor
@@ -103,7 +106,7 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     from pyro.infer.autoguide import init_to_mean
 
     # region prior sampling
-    prior = gpp.run_predictive(
+    prior_inference = gpp.run_predictive(
         prob_model=prob_model,
         geo_model=geo_model,
         y_obs_list=y_obs_list,
@@ -114,30 +117,51 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     # endregion
 
     # region inference
-    pyro.primitives.enable_validation(is_validate=True)
-    nuts_kernel = NUTS(
-        prob_model,
-        step_size=0.0085,
-        adapt_step_size=True,
-        target_accept_prob=0.9,
-        max_tree_depth=10,
-        init_strategy=init_to_mean
-    )
-    mcmc = MCMC(
-        kernel=nuts_kernel,
-        num_samples=200,
-        warmup_steps=50,
-        disable_validation=False
-    )
-    mcmc.run(geo_model, y_obs_list)
-    posterior_samples = mcmc.get_samples()
-    posterior_predictive_fn = Predictive(
-        model=prob_model,
-        posterior_samples=posterior_samples
+    data = gpp.run_nuts_inference(
+        prob_model=prob_model,
+        geo_model=geo_model,
+        y_obs_list=y_obs_list,
+        config=NUTSConfig(
+            step_size=0.0085,
+            adapt_step_size=True,
+            target_accept_prob=0.9,
+            max_tree_depth=10,
+            init_strategy='auto',
+            num_samples=200,
+            warmup_steps=50,
+        ),
+        plot_trace=False,
+        run_posterior_predictive=True
     )
-    posterior_predictive = posterior_predictive_fn(geo_model, y_obs_list)
-
-    data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)
+    # pyro.primitives.enable_validation(is_validate=True)
+    # nuts_kernel = NUTS(
+    #     prob_model,
+    #     step_size=0.0085,
+    #     adapt_step_size=True,
+    #     target_accept_prob=0.9,
+    #     max_tree_depth=10,
+    #     init_strategy=init_to_mean
+    # )
+    # mcmc = MCMC(
+    #     kernel=nuts_kernel,
+    #     num_samples=200,
+    #     warmup_steps=50,
+    #     disable_validation=False
+    # )
+    # mcmc.run(geo_model, y_obs_list)
+    # posterior_samples = mcmc.get_samples()
+    # posterior_predictive_fn = Predictive(
+    #     model=prob_model,
+    #     posterior_samples=posterior_samples
+    # )
+    # posterior_predictive = posterior_predictive_fn(geo_model, y_obs_list)
+    # 
+    # data = az.from_pyro(
+    #     posterior=mcmc,
+    #     # prior=prior.to_dict()["prior"],
+    #     posterior_predictive=posterior_predictive
+    # )
+    data.extend(prior_inference)
     # endregion
 
     # print("Number of interpolations: ", geo_model.counter)
@@ -149,7 +173,7 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     plt.show()
     from gempy_probability.modules.plot.plot_posterior import default_red, default_blue
     az.plot_density(
-        data=[data.posterior_predictive, data.prior_predictive],
+        data=[data.posterior_predictive, data.prior],
         shade=.9,
         var_names=[r'$\mu_{thickness}$'],
         data_labels=["Posterior Predictive", "Prior Predictive"],
@@ -165,6 +189,15 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
     )
     plt.show()
 
+    az.plot_density(
+        data=[data.posterior_predictive, data.prior],
+        shade=.9,
+        hdi_prob=.99,
+        data_labels=["Posterior", "Prior"],
+        colors=[default_red, default_blue],
+    )
+    plt.show()
+
     # Test posterior mean values
     posterior_top_mean = float(data.posterior[r'$\mu_{top}$'].mean())
     target_top = 0.00875
@@ -180,5 +213,3 @@ 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())}")
-
-

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-from ._pyro_runner import run_predictive`
	`1`	`+from ._pyro_runner import run_predictive, run_nuts_inference, run_mcmc_for_NUTS`