Add probabilistic modeling API with model factory and runner (#11)

# Probabilistic Modeling Framework for GemPy

This PR introduces a comprehensive framework for probabilistic modeling in GemPy, making it easier to define, run, and analyze probabilistic geological models.

## Key Features

- Added a factory function `make_gempy_pyro_model` to create Pyro models for GemPy with clean separation of concerns
- Introduced `GemPyPyroModel` type alias for better type clarity throughout the codebase
- Created standardized API for running probabilistic models with `run_predictive`, `run_nuts_inference`, and `run_mcmc_for_NUTS`
- Implemented a configurable NUTS sampler with `NUTSConfig` dataclass
- Reorganized likelihood functions into a more modular structure
- Added comprehensive documentation for model creation and usage

## Implementation Details

- Moved backend configuration to `run_gempy_forward` for consistent execution
- Refactored likelihood modules for better organization and reusability
- Added utility functions for common probabilistic modeling tasks
- Created test cases demonstrating the new framework's capabilities
- Removed outdated gravity inversion example in favor of the new approach

This PR lays the groundwork for more sophisticated probabilistic modeling in GemPy, making it easier for users to define custom models while maintaining a clean separation between model definition, execution, and analysis.

Author: Leguark

docs/dev_logs/2025_May.md

@@ -46,12 +46,15 @@
 ### Likelihood functions
 - ? Should I start to create a catalog of likelihood functions that goes with gempy?
   - This would be useful to be able to visualize whatever the likelihood function represents independent of the probabilistic model.
-  - 
+
 ### What would I want to have to do StonePark example as flowless as possible? 
 1. [x] A way to save the model already with the correct **nugget effects** instead of having to "load nuggets" (TODO in gempy)
 2. [ ] Likelihood function plug and play
 3. [ ] Some sort of pre-analysis if the model looks good. Maybe some metrics about how long it takes to run the model. 
    1. Condition number
    2. Acceptance rate of the posterior
 4. Easy pykeops setting using .env
-5. Some sort of easy set-up for having a reduced number of priors
+5. Some sort of easy set-up for having a reduced number of priors
+
+
+## Probabilistic Modeling definition

examples/examples/gravity_inversion_I.py

@@ -1 +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, run_mcmc_for_NUTS, run_nuts_inference
+
 from ._version import __version__

gempy_probability/__init__.py

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

gempy_probability/api/model_runner/__init__.py

@@ -0,0 +1,90 @@
+import arviz as az
+import matplotlib.pyplot as plt
+import torch
+from pyro.infer import MCMC
+from pyro.infer import NUTS
+from pyro.infer import Predictive
+from pyro.infer.autoguide import init_to_mean
+
+import gempy as gp
+import gempy_probability as gpp
+from ...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:
+    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
+
+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

gempy_probability/api/model_runner/_pyro_runner.py

@@ -0,0 +1,12 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class NUTSConfig:
+    step_size: float = 0.0085
+    adapt_step_size: bool = True
+    target_accept_prob: float = 0.9
+    max_tree_depth: int = 10
+    init_strategy: str = 'auto'
+    num_samples: int = 200 
+    warmup_steps: int = 50

gempy_probability/core/samplers_data.py

@@ -0,0 +1 @@
+from ._gempy_fw import run_gempy_forward

gempy_probability/modules/forwards/__init__.py

@@ -0,0 +1,19 @@
+import gempy as gp
+import gempy_engine
+from gempy_engine.core.backend_tensor import BackendTensor
+from gempy_engine.core.data.interpolation_input import InterpolationInput
+
+
+def run_gempy_forward(interp_input: InterpolationInput, geo_model: gp.data.GeoModel) -> gp.data.Solutions:
+    
+    # compute model
+    sol = gempy_engine.compute_model(
+        interpolation_input=interp_input,
+        options=geo_model.interpolation_options,
+        data_descriptor=geo_model.input_data_descriptor,
+        geophysics_input=geo_model.geophysics_input
+    )
+    
+    return sol
+
+

gempy_probability/modules/forwards/_gempy_fw.py

@@ -1,3 +1,2 @@
-from ._apparent_thickness import apparent_thickness_likelihood, apparent_thickness_likelihood_II
-from ._gravity_inv import gravity_inversion_likelihood
+from ._likelihood_functions import thickness_likelihood