From 65021dfe38c1ca65ca0fc6f7c8cd4e3d5132160c Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 12:17:06 +0100 Subject: [PATCH 01/14] file --- examples/examples/gravity_inversion_I.py | 71 ------------------- .../model_definition/prob_model_factory.py | 0 2 files changed, 71 deletions(-) delete mode 100644 examples/examples/gravity_inversion_I.py create mode 100644 gempy_probability/modules/model_definition/prob_model_factory.py diff --git a/examples/examples/gravity_inversion_I.py b/examples/examples/gravity_inversion_I.py deleted file mode 100644 index c45c3b3..0000000 --- a/examples/examples/gravity_inversion_I.py +++ /dev/null @@ -1,71 +0,0 @@ -""" -Probabilistic Inversion Example: Geological Model: Making the model stable --------------------------------------------------------------------------- - - -""" - - -import os -import time - -import arviz as az -import numpy as np -import pyro -import pyro.distributions as dist -import torch -import xarray as xr -from dotenv import dotenv_values -from matplotlib import pyplot as plt -from pyro.infer import MCMC, NUTS, Predictive - -import gempy as gp -import gempy_engine -import gempy_viewer as gpv - -from examples.examples._aux_func_gravity_inversion import calculate_scale_shift, gaussian_kernel, initialize_geo_model, setup_geophysics -from examples.examples._aux_func_gravity_inversion_II import process_file - -from gempy_engine.core.backend_tensor import BackendTensor -from gempy_probability.modules.plot.plot_posterior import default_red, default_blue - -# %% -# Config -seed = 123456 -torch.manual_seed(seed) -pyro.set_rng_seed(seed) - -# %% -# Start the timer for benchmarking purposes -start_time = time.time() - -# %% -# Load necessary configuration and paths from environment variables -config = dotenv_values() -path = config.get("PATH_TO_MODEL_1_Subsurface") - -# %% -# Initialize lists to store structural elements for the geological model -structural_elements = [] -global_extent = None -color_gen = gp.data.ColorsGenerator() - -# %% -# Process each .nc file in the specified directory for model construction -for filename in os.listdir(path): - base, ext = os.path.splitext(filename) - if ext == '.nc': - structural_element, global_extent = process_file(os.path.join(path, filename), global_extent, color_gen) - structural_elements.append(structural_element) - -# %% -# Configure GemPy for geological modeling with PyTorch backend -BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH, dtype="float64") - -geo_model = initialize_geo_model( - structural_elements=structural_elements, - extent=(np.array(global_extent)), - topography=(xr.open_dataset(os.path.join(path, "Topography.nc"))), - load_nuggets=False -) - diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py new file mode 100644 index 0000000..e69de29 From 99b494bf037b9d817336331d3ecb5e0d5383bffa Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 12:57:59 +0100 Subject: [PATCH 02/14] [WIP] Developing interface for probabilistic models --- docs/dev_logs/2025_May.md | 7 +- .../model_definition/model_examples.py | 59 ++++++- .../model_definition/prob_model_factory.py | 32 ++++ tests/test_prob_model/test_prob_factory.py | 154 ++++++++++++++++++ 4 files changed, 247 insertions(+), 5 deletions(-) create mode 100644 tests/test_prob_model/test_prob_factory.py diff --git a/docs/dev_logs/2025_May.md b/docs/dev_logs/2025_May.md index 8858957..eb3e7af 100644 --- a/docs/dev_logs/2025_May.md +++ b/docs/dev_logs/2025_May.md @@ -46,7 +46,7 @@ ### 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 @@ -54,4 +54,7 @@ 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 \ No newline at end of file +5. Some sort of easy set-up for having a reduced number of priors + + +## Probabilistic Modeling definition \ No newline at end of file diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index 29119c4..8e42822 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -1,13 +1,14 @@ import gempy as gp import gempy.core.data import gempy_engine -from gempy.modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame -from gempy_probability.modules.likelihoods import apparent_thickness_likelihood +from gempy.modules.data_manipulation import interpolation_input_from_structural_frame +from ..likelihoods import apparent_thickness_likelihood import pyro -import pyro.distributions as dist import torch +from gempy_probability.modules.model_definition.prob_model_factory import make_pyro_model + def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_list): """ @@ -72,3 +73,55 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li ) + +import pyro.distributions as dist + +# 1) Define your priors +gravity_priors = { + # e.g. density contrast of layer index 3 + "mu_density": dist.Normal(2.62, 0.5) +} + +# 2) Wrap your forward pass into a small fn +def run_gempy_forward(samples, geo_model): + + mu_top = samples["mu_density"] + + interp_input = interpolation_input_from_structural_frame(geo_model) + + indices__ = (torch.tensor([0]), torch.tensor([2])) # * This has to be Tensors + new_tensor: torch.Tensor = torch.index_put( + input=interp_input.surface_points.sp_coords, + indices=indices__, + values=mu_top + ) + interp_input.surface_points.sp_coords = new_tensor + + # 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 + ) + + thickness = apparent_thickness_likelihood(sol) + return thickness + +# 3) Define your likelihood factory +def gravity_likelihood(simulated_gravity): + # e.g. multivariate normal with precomputed cov matrix + return dist.MultivariateNormal(simulated_gravity, covariance_matrix) + +def thickness_likelihood(simulated_thickness): + # e.g. multivariate normal with precomputed cov matrix + return dist.Normal(simulated_thickness, 25) + +# 4) Build the model +pyro_gravity_model = make_pyro_model( + priors=gravity_priors, + forward_fn=run_gempy_forward, + likelihood_fn=thickness_likelihood, + obs_name="obs_gravity" +) + diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py index e69de29..1ac6bf0 100644 --- a/gempy_probability/modules/model_definition/prob_model_factory.py +++ b/gempy_probability/modules/model_definition/prob_model_factory.py @@ -0,0 +1,32 @@ +import pyro +import torch +from typing import Callable, Dict +from pyro.distributions import Distribution + +def make_pyro_model( + priors: Dict[str, Distribution], + forward_fn: Callable[..., torch.Tensor], + likelihood_fn: Callable[[torch.Tensor], Distribution], + obs_name: str = "obs" +): + """ + Build a Pyro model function with: + - priors: map parameter names -> pyro.distributions objects + - forward_fn: fn(samples_dict, *args)-> simulated quantity (torch.Tensor) + - likelihood_fn: fn(simulated)-> a Pyro Distribution over data + - obs_name: name of the observed site + """ + def model(geo_model, obs_data): + # 1) Sample each prior + samples = {} + for name, dist in priors.items(): + samples[name] = pyro.sample(name, dist) + + # 2) Run your forward geological model + simulated = forward_fn(samples, geo_model) + + # 3) Build likelihood and observe + lik_dist = likelihood_fn(simulated) + pyro.sample(obs_name, lik_dist, obs=obs_data) + + return model \ No newline at end of file diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py new file mode 100644 index 0000000..bac3791 --- /dev/null +++ b/tests/test_prob_model/test_prob_factory.py @@ -0,0 +1,154 @@ +import numpy as np +import os +import pyro.distributions as dist +import torch + +import gempy as gp +from gempy_engine.core.backend_tensor import BackendTensor + +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')) + geo_model = gp.create_geomodel( + project_name='Wells', + extent=[0, 12000, -500, 500, 0, 4000], + refinement=3, + importer_helper=gp.data.ImporterHelper( + path_to_orientations=os.path.join(data_path, "2-layers", "2-layers_orientations.csv"), + path_to_surface_points=os.path.join(data_path, "2-layers", "2-layers_surface_points.csv") + ) + ) + + # TODO: Make this a more elegant way + BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) + + # 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. + + # region Minimal grid for the specific likelihood function + x_loc = 6000 + y_loc = 0 + z_loc = np.linspace(0, 4000, 100) + xyz_coord = np.array([[x_loc, y_loc, z] for z in z_loc]) + gp.set_custom_grid(geo_model.grid, xyz_coord=xyz_coord) + # endregion + + geo_model.grid.active_grids = gp.data.Grid.GridTypes.CUSTOM + assert geo_model.grid.values.shape[0] == 100, "Custom grid should have 100 cells" + # geo_model.counter = 0 + gp.compute_model( + gempy_model=geo_model, + validate_serialization=False + ) + BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) + + normal = dist.Normal( + loc=(geo_model.surface_points_copy_transformed.xyz[0, 2]), + scale=torch.tensor(0.1, dtype=torch.float64) + ) + + from gempy_probability.modules.model_definition.model_examples import two_wells_prob_model_I, pyro_gravity_model + + _prob_run( + geo_model=geo_model, + prob_model=pyro_gravity_model, + normal=normal, + y_obs_list=torch.tensor([200, 210, 190]) + ) + + +def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, + normal: torch.distributions.Distribution, y_obs_list: torch.Tensor) -> None: + # Run prior sampling and visualization + 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 + + # region prior sampling + predictive = Predictive( + model=prob_model, + num_samples=50 + ) + 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 + + # 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, normal, 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, normal, y_obs_list) + + data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive) + # endregion + + # print("Number of interpolations: ", geo_model.counter) + + if True: # * Save the arviz data + data.to_netcdf("arviz_data.nc") + + az.plot_trace(data) + plt.show() + from gempy_probability.modules.plot.plot_posterior import default_red, default_blue + az.plot_density( + data=[data.posterior_predictive, data.prior_predictive], + shade=.9, + var_names=[r'$\mu_{thickness}$'], + data_labels=["Posterior Predictive", "Prior Predictive"], + colors=[default_red, default_blue], + ) + plt.show() + az.plot_density( + data=[data, 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 + target_thickness = 223 + assert abs(posterior_top_mean - target_top) < 0.0200, f"Top layer mean {posterior_top_mean} outside expected range" + posterior_thickness_mean = float(data.posterior_predictive[r'$\mu_{thickness}$'].mean()) + assert abs(posterior_thickness_mean - target_thickness) < 5, f"Thickness mean {posterior_thickness_mean} outside expected range" + # Test convergence diagnostics + assert float(data.sample_stats.diverging.sum()) == 0, "MCMC sampling has divergences" + + print("Posterior mean values:") + print(f"Top layer mean: {posterior_top_mean}") + print(f"Thickness mean: {posterior_thickness_mean}") + print("MCMC convergence diagnostics:") + print(f"Divergences: {float(data.sample_stats.diverging.sum())}") From c41b2c2c8f13d688f0471b7d98a2169f95719527 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 13:12:46 +0100 Subject: [PATCH 03/14] [WIP] Getting there --- .../modules/model_definition/model_examples.py | 8 ++++++-- tests/test_prob_model/test_prob_factory.py | 6 +++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index 8e42822..77afea0 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -79,13 +79,17 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li # 1) Define your priors gravity_priors = { # e.g. density contrast of layer index 3 - "mu_density": dist.Normal(2.62, 0.5) + # "mu_density": dist.Normal(2.62, 0.5), + "mu_thickness": dist.Normal( + loc=torch.tensor(0.0487, dtype=torch.float64), + scale=torch.tensor(0.1, dtype=torch.float64) + ) } # 2) Wrap your forward pass into a small fn def run_gempy_forward(samples, geo_model): - mu_top = samples["mu_density"] + mu_top = samples["mu_thickness"] interp_input = interpolation_input_from_structural_frame(geo_model) diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index bac3791..dd4138d 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -101,20 +101,20 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, warmup_steps=50, disable_validation=False ) - mcmc.run(geo_model, normal, y_obs_list) + 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, normal, y_obs_list) + posterior_predictive = posterior_predictive_fn(geo_model, y_obs_list) data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive) # endregion # print("Number of interpolations: ", geo_model.counter) - if True: # * Save the arviz data + if False: # * Save the arviz data data.to_netcdf("arviz_data.nc") az.plot_trace(data) From f460733c45a7a4475bbce6799cc0b97f0eddc173 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 13:16:08 +0100 Subject: [PATCH 04/14] [WIP/TEST] prob_factory test running for the first time --- gempy_probability/modules/model_definition/model_examples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index 77afea0..e0e11bb 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -80,7 +80,7 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li gravity_priors = { # e.g. density contrast of layer index 3 # "mu_density": dist.Normal(2.62, 0.5), - "mu_thickness": dist.Normal( + r'$\mu_{top}$': dist.Normal( loc=torch.tensor(0.0487, dtype=torch.float64), scale=torch.tensor(0.1, dtype=torch.float64) ) @@ -89,7 +89,7 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li # 2) Wrap your forward pass into a small fn def run_gempy_forward(samples, geo_model): - mu_top = samples["mu_thickness"] + mu_top = samples["$\mu_{top}$"] interp_input = interpolation_input_from_structural_frame(geo_model) From ab4be33c3b14cc8e90d53a3133ee57ebb10dae3b Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 13:44:37 +0100 Subject: [PATCH 05/14] [TEST] Run old test --- .../modules/model_definition/model_examples.py | 2 +- tests/test_prob_model/test_prob_I.py | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index e0e11bb..f6a998d 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -45,7 +45,7 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li # region Forward model computation - geo_model.counter +=1 + # geo_model.counter +=1 # * Compute the geological model geo_model.solutions = gempy_engine.compute_model( diff --git a/tests/test_prob_model/test_prob_I.py b/tests/test_prob_model/test_prob_I.py index 5304c3d..f3ac558 100644 --- a/tests/test_prob_model/test_prob_I.py +++ b/tests/test_prob_model/test_prob_I.py @@ -38,8 +38,7 @@ def test_basic_gempy_I() -> None: geo_model.grid.active_grids = gp.data.Grid.GridTypes.CUSTOM assert geo_model.grid.values.shape[0] == 100, "Custom grid should have 100 cells" - geo_model.counter = 0 - gp.compute_model(gempy_model=geo_model) + gp.compute_model(gempy_model=geo_model, validate_serialization=False) BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) normal = dist.Normal( @@ -74,7 +73,7 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, num_samples=50 ) prior = predictive(geo_model, normal, y_obs_list) - print("Number of interpolations: ", geo_model.counter) + # print("Number of interpolations: ", geo_model.counter) data = az.from_pyro(prior=prior) az.plot_trace(data.prior) @@ -109,7 +108,7 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive) # endregion - print("Number of interpolations: ", geo_model.counter) + # print("Number of interpolations: ", geo_model.counter) if True: # * Save the arviz data data.to_netcdf("arviz_data.nc") From 94f0548e1dde55c62a9c6db4b461a5e1c14cde2e Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 14:26:49 +0100 Subject: [PATCH 06/14] [CLN/ENH] Better definition for specific gempy_pyro models --- gempy_probability/__init__.py | 2 + .../modules/forwards/__init__.py | 1 + .../modules/forwards/_gempy_fw.py | 22 +++++++ .../likelihoods/_likelihood_functions.py | 15 +++++ .../model_definition/model_examples.py | 58 ------------------- .../model_definition/prob_model_factory.py | 34 ++++++++++- tests/test_prob_model/test_prob_factory.py | 47 ++++++++++++--- 7 files changed, 112 insertions(+), 67 deletions(-) create mode 100644 gempy_probability/modules/forwards/__init__.py create mode 100644 gempy_probability/modules/forwards/_gempy_fw.py create mode 100644 gempy_probability/modules/likelihoods/_likelihood_functions.py diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py index 8dee4bf..7c4dd8a 100644 --- a/gempy_probability/__init__.py +++ b/gempy_probability/__init__.py @@ -1 +1,3 @@ +from .modules.model_definition.prob_model_factory import make_gempy_pyro_model + from ._version import __version__ diff --git a/gempy_probability/modules/forwards/__init__.py b/gempy_probability/modules/forwards/__init__.py new file mode 100644 index 0000000..09b2990 --- /dev/null +++ b/gempy_probability/modules/forwards/__init__.py @@ -0,0 +1 @@ +from ._gempy_fw import run_gempy_forward \ No newline at end of file diff --git a/gempy_probability/modules/forwards/_gempy_fw.py b/gempy_probability/modules/forwards/_gempy_fw.py new file mode 100644 index 0000000..b55f7f4 --- /dev/null +++ b/gempy_probability/modules/forwards/_gempy_fw.py @@ -0,0 +1,22 @@ +import torch +from pyro.distributions import Distribution + +import gempy_engine +from gempy.modules.data_manipulation import interpolation_input_from_structural_frame +import gempy as gp +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 + + diff --git a/gempy_probability/modules/likelihoods/_likelihood_functions.py b/gempy_probability/modules/likelihoods/_likelihood_functions.py new file mode 100644 index 0000000..f5be1a3 --- /dev/null +++ b/gempy_probability/modules/likelihoods/_likelihood_functions.py @@ -0,0 +1,15 @@ +import torch +import pyro.distributions as dist +import gempy as gp +from . import apparent_thickness_likelihood + + +def gravity_likelihood(geo_model: gp.data.Solutions): + # e.g. multivariate normal with precomputed cov matrix + raise NotImplementedError("This function is not yet implemented") + return dist.MultivariateNormal(simulated_gravity, covariance_matrix) + +def thickness_likelihood(solutions: gp.data.Solutions) -> dist: + # e.g. multivariate normal with precomputed cov matrix + simulated_thickness = apparent_thickness_likelihood(solutions) + return dist.Normal(simulated_thickness, 25) diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index f6a998d..ec9d6a9 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -7,8 +7,6 @@ import pyro import torch -from gempy_probability.modules.model_definition.prob_model_factory import make_pyro_model - def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_list): """ @@ -73,59 +71,3 @@ def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_li ) - -import pyro.distributions as dist - -# 1) Define your priors -gravity_priors = { - # e.g. density contrast of layer index 3 - # "mu_density": dist.Normal(2.62, 0.5), - r'$\mu_{top}$': dist.Normal( - loc=torch.tensor(0.0487, dtype=torch.float64), - scale=torch.tensor(0.1, dtype=torch.float64) - ) -} - -# 2) Wrap your forward pass into a small fn -def run_gempy_forward(samples, geo_model): - - mu_top = samples["$\mu_{top}$"] - - interp_input = interpolation_input_from_structural_frame(geo_model) - - indices__ = (torch.tensor([0]), torch.tensor([2])) # * This has to be Tensors - new_tensor: torch.Tensor = torch.index_put( - input=interp_input.surface_points.sp_coords, - indices=indices__, - values=mu_top - ) - interp_input.surface_points.sp_coords = new_tensor - - # 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 - ) - - thickness = apparent_thickness_likelihood(sol) - return thickness - -# 3) Define your likelihood factory -def gravity_likelihood(simulated_gravity): - # e.g. multivariate normal with precomputed cov matrix - return dist.MultivariateNormal(simulated_gravity, covariance_matrix) - -def thickness_likelihood(simulated_thickness): - # e.g. multivariate normal with precomputed cov matrix - return dist.Normal(simulated_thickness, 25) - -# 4) Build the model -pyro_gravity_model = make_pyro_model( - priors=gravity_priors, - forward_fn=run_gempy_forward, - likelihood_fn=thickness_likelihood, - obs_name="obs_gravity" -) - diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py index 1ac6bf0..1e0c985 100644 --- a/gempy_probability/modules/model_definition/prob_model_factory.py +++ b/gempy_probability/modules/model_definition/prob_model_factory.py @@ -3,7 +3,37 @@ from typing import Callable, Dict from pyro.distributions import Distribution -def make_pyro_model( +from gempy_probability.modules.forwards import run_gempy_forward +import gempy as gp + + +def make_gempy_pyro_model( + priors: Dict[str, Distribution], + set_interp_input_fn: Callable[..., gp.data.Solutions], + likelihood_fn: Callable[[gp.data.Solutions], Distribution], + obs_name: str = "obs" +): + + def model(geo_model, obs_data): + # 1) Sample each prior + samples = {} + for name, dist in priors.items(): + samples[name] = pyro.sample(name, dist) + + # 3) Run your forward geological model + simulated: gp.data.Solutions = run_gempy_forward( + interp_input=set_interp_input_fn(samples, geo_model), + geo_model=geo_model + ) + + # 4) Build likelihood and observe + lik_dist = likelihood_fn(simulated) + pyro.sample(obs_name, lik_dist, obs=obs_data) + + return model + + +def make_generic_pyro_model( priors: Dict[str, Distribution], forward_fn: Callable[..., torch.Tensor], likelihood_fn: Callable[[torch.Tensor], Distribution], @@ -29,4 +59,4 @@ def model(geo_model, obs_data): lik_dist = likelihood_fn(simulated) pyro.sample(obs_name, lik_dist, obs=obs_data) - return model \ No newline at end of file + return model diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index dd4138d..26cced5 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -1,10 +1,18 @@ import numpy as np import os +import pyro import pyro.distributions as dist import torch +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 + +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__)) @@ -44,23 +52,48 @@ def test_prob_model_factory() -> None: ) BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) - normal = dist.Normal( - loc=(geo_model.surface_points_copy_transformed.xyz[0, 2]), - scale=torch.tensor(0.1, dtype=torch.float64) + # 1) Define your priors + model_priors = { + r'$\mu_{top}$': dist.Normal( + loc=geo_model.surface_points_copy_transformed.xyz[0, 2], + scale=torch.tensor(0.1, dtype=torch.float64) + ) + } + + from gempy_probability.modules.likelihoods._likelihood_functions import thickness_likelihood + + pyro_gravity_model = gpp.make_gempy_pyro_model( + priors=model_priors, + set_interp_input_fn=modify_z_for_surface_point1, + likelihood_fn=thickness_likelihood, + obs_name="obs_gravity" ) - from gempy_probability.modules.model_definition.model_examples import two_wells_prob_model_I, pyro_gravity_model - _prob_run( geo_model=geo_model, prob_model=pyro_gravity_model, - normal=normal, y_obs_list=torch.tensor([200, 210, 190]) ) + +def modify_z_for_surface_point1( + samples: dict[str, Distribution], + geo_model: gp.data.GeoModel, +) -> InterpolationInput: + 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 + values=(samples[prior_key]) + ) + interp_input.surface_points.sp_coords = new_tensor + return interp_input def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, - normal: torch.distributions.Distribution, y_obs_list: torch.Tensor) -> None: + y_obs_list: torch.Tensor) -> None: # Run prior sampling and visualization from pyro.infer import Predictive import pyro From a464d5aab5f36b28a9b6922f7de06c82dc8dabe2 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 14:41:10 +0100 Subject: [PATCH 07/14] [CLN] Define better likelihood module --- gempy_probability/__init__.py | 1 + gempy_probability/modules/likelihoods/__init__.py | 3 +-- gempy_probability/modules/likelihoods/_gravity_inv.py | 9 --------- .../modules/likelihoods/_likelihood_functions.py | 3 +-- .../modules/likelihoods/gravity_likelihoods/__init__.py | 1 - .../modules/model_definition/model_examples.py | 7 ++++--- gempy_probability/modules/utils/__init__.py | 0 .../{likelihoods/gravity_likelihoods => utils}/_utils.py | 0 tests/test_likelihoods/test_geometric_likelihoods.py | 2 +- tests/test_prob_model/test_prob_factory.py | 6 +++--- 10 files changed, 11 insertions(+), 21 deletions(-) delete mode 100644 gempy_probability/modules/likelihoods/_gravity_inv.py delete mode 100644 gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.py create mode 100644 gempy_probability/modules/utils/__init__.py rename gempy_probability/modules/{likelihoods/gravity_likelihoods => utils}/_utils.py (100%) diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py index 7c4dd8a..7ffc48e 100644 --- a/gempy_probability/__init__.py +++ b/gempy_probability/__init__.py @@ -1,3 +1,4 @@ from .modules.model_definition.prob_model_factory import make_gempy_pyro_model +from .modules import likelihoods from ._version import __version__ diff --git a/gempy_probability/modules/likelihoods/__init__.py b/gempy_probability/modules/likelihoods/__init__.py index e88d21a..0f9ec2d 100644 --- a/gempy_probability/modules/likelihoods/__init__.py +++ b/gempy_probability/modules/likelihoods/__init__.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 diff --git a/gempy_probability/modules/likelihoods/_gravity_inv.py b/gempy_probability/modules/likelihoods/_gravity_inv.py deleted file mode 100644 index 889af2d..0000000 --- a/gempy_probability/modules/likelihoods/_gravity_inv.py +++ /dev/null @@ -1,9 +0,0 @@ -import pyro -import torch -import gempy as gp - -def gravity_inversion_likelihood(model_solutions: gp.data.Solutions) -> torch.Tensor: - simulated_geophysics = model_solutions.gravity - pyro.deterministic(r'$\mu_{gravity}$', simulated_geophysics) - - return simulated_geophysics \ No newline at end of file diff --git a/gempy_probability/modules/likelihoods/_likelihood_functions.py b/gempy_probability/modules/likelihoods/_likelihood_functions.py index f5be1a3..d053270 100644 --- a/gempy_probability/modules/likelihoods/_likelihood_functions.py +++ b/gempy_probability/modules/likelihoods/_likelihood_functions.py @@ -1,7 +1,6 @@ -import torch import pyro.distributions as dist import gempy as gp -from . import apparent_thickness_likelihood +from ._apparent_thickness import apparent_thickness_likelihood def gravity_likelihood(geo_model: gp.data.Solutions): diff --git a/gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.py b/gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.py deleted file mode 100644 index 16c221b..0000000 --- a/gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from ._utils import calculate_scale_shift, gaussian_kernel \ No newline at end of file diff --git a/gempy_probability/modules/model_definition/model_examples.py b/gempy_probability/modules/model_definition/model_examples.py index ec9d6a9..608455d 100644 --- a/gempy_probability/modules/model_definition/model_examples.py +++ b/gempy_probability/modules/model_definition/model_examples.py @@ -1,14 +1,15 @@ import gempy as gp -import gempy.core.data import gempy_engine from gempy.modules.data_manipulation import interpolation_input_from_structural_frame -from ..likelihoods import apparent_thickness_likelihood import pyro +from pyro import distributions as dist import torch +from gempy_probability.modules.likelihoods._apparent_thickness import apparent_thickness_likelihood -def two_wells_prob_model_I(geo_model: gempy.core.data.GeoModel, normal, y_obs_list): + +def two_wells_prob_model_I(geo_model: gp.data.GeoModel, normal, y_obs_list): """ This Pyro model represents the probabilistic aspects of the geological model. It defines a prior distribution for the top layer's location and diff --git a/gempy_probability/modules/utils/__init__.py b/gempy_probability/modules/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/gempy_probability/modules/likelihoods/gravity_likelihoods/_utils.py b/gempy_probability/modules/utils/_utils.py similarity index 100% rename from gempy_probability/modules/likelihoods/gravity_likelihoods/_utils.py rename to gempy_probability/modules/utils/_utils.py diff --git a/tests/test_likelihoods/test_geometric_likelihoods.py b/tests/test_likelihoods/test_geometric_likelihoods.py index d6ee9ab..1b2e7c5 100644 --- a/tests/test_likelihoods/test_geometric_likelihoods.py +++ b/tests/test_likelihoods/test_geometric_likelihoods.py @@ -5,7 +5,7 @@ import gempy_viewer as gpv from gempy.core.data.grid_modules import CustomGrid from gempy_engine.core.backend_tensor import BackendTensor -from gempy_probability.modules.likelihoods import apparent_thickness_likelihood_II +from gempy_probability.modules.likelihoods._apparent_thickness import apparent_thickness_likelihood_II def test_gempy_posterior_I(): diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index 26cced5..4781718 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -60,12 +60,10 @@ def test_prob_model_factory() -> None: ) } - from gempy_probability.modules.likelihoods._likelihood_functions import thickness_likelihood - pyro_gravity_model = gpp.make_gempy_pyro_model( priors=model_priors, set_interp_input_fn=modify_z_for_surface_point1, - likelihood_fn=thickness_likelihood, + likelihood_fn=gpp.likelihoods.thickness_likelihood, obs_name="obs_gravity" ) @@ -79,6 +77,8 @@ 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 From be25c11854c2ddd20b2f907b17b91d548517f098 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 14:56:19 +0100 Subject: [PATCH 08/14] [DOC] Adding a ton of docs to construct the model --- .../model_definition/prob_model_factory.py | 98 +++++++++++++++++-- 1 file changed, 90 insertions(+), 8 deletions(-) diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py index 1e0c985..ea6360d 100644 --- a/gempy_probability/modules/model_definition/prob_model_factory.py +++ b/gempy_probability/modules/model_definition/prob_model_factory.py @@ -9,30 +9,112 @@ def make_gempy_pyro_model( priors: Dict[str, Distribution], - set_interp_input_fn: Callable[..., gp.data.Solutions], + set_interp_input_fn: Callable[ + [Dict[str, torch.Tensor], gp.data.GeoModel], + gp.data.InterpolationInput + ], likelihood_fn: Callable[[gp.data.Solutions], Distribution], obs_name: str = "obs" -): +) -> Callable[[gp.data.GeoModel, torch.Tensor], None]: + """ + Factory to produce a Pyro model for GemPy forward simulations. - def model(geo_model, obs_data): - # 1) Sample each prior - samples = {} + This returns a `model(geo_model, obs_data)` function that you can hand + directly to Pyro's inference APIs (`Predictive`, `MCMC`, …). Internally + the generated model does: + + 1. Samples each key/value in `priors` via `pyro.sample(name, dist)`. + 2. Calls `set_interp_input_fn(samples, geo_model)` to inject those samples + into a new `InterpolationInput`. + 3. Runs the GemPy forward solver with `run_gempy_forward(...)`. + 4. Wraps the resulting `Solutions` in `likelihood_fn(...)` to get a Pyro + distribution, then observes `obs_data` under that distribution. + + Parameters + ---------- + priors : Dict[str, Distribution] + A mapping from Pyro sample‐site names to Pyro Distribution objects + defining your prior beliefs over each parameter. + set_interp_input_fn : Callable[[samples, geo_model], InterpolationInput] + A user function which receives: + + * `samples` (Dict[str, Tensor]): the values drawn from your priors + * `geo_model` (gempy.core.data.GeoModel): the base geological model + + and must return a ready‐to‐use `InterpolationInput` for GemPy. + likelihood_fn : Callable[[Solutions], Distribution] + A function mapping the GemPy forward‐model output (`Solutions`) + to a Pyro Distribution representing the likelihood of the observed data. + obs_name : str, optional + The Pyro site name under which `obs_data` will be observed. + Defaults to `"obs"`. + + Returns + ------- + model : Callable[[GeoModel, Tensor], None] + A Pyro‐compatible model taking: + + * `geo_model`: your GemPy GeoModel object + * `obs_data`: a torch.Tensor of observations + + This function has no return value; it registers its random draws + via `pyro.sample`. + + Example + ------- + >>> import pyro.distributions as dist + >>> from gempy_probability import make_gempy_pyro_model + >>> from gempy.modules.data_manipulation import interpolation_input_from_structural_frame + >>> import gempy_probability + >>> + >>> # 1) Define a simple Normal prior on μ + >>> priors = {"μ": dist.Normal(0., 1.)} + >>> + >>> # 2) Function to inject μ into your interp-input + >>> def set_input(samples, gm): + ... inp = interpolation_input_from_structural_frame(gm) + ... inp.surface_points.sp_coords = torch.index_put( + ... input=inp.surface_points.sp_coords, + ... indices=(torch.tensor([0]), torch.tensor([2])), # * This has to be Tensors + ... values=(samples["μ"]) + ... ) + ... return inp + >>> + >>> + >>> # 4) Build the model + >>> pyro_model = make_gempy_pyro_model( + ... priors=priors, + ... set_interp_input_fn=set_input, + ... likelihood_fn=gempy_probability.likelihoods.thickness_likelihood, + ... obs_name="y") + >>> + >>> + >>> # Now this can be used with Predictive or MCMC directly: + >>> # Predictive(pyro_model, num_samples=100)(geo_model, obs_tensor) + """ + def model(geo_model: gp.data.GeoModel, obs_data: torch.Tensor): + # 1) Sample from the user‐supplied priors + samples: Dict[str, torch.Tensor] = {} for name, dist in priors.items(): samples[name] = pyro.sample(name, dist) - # 3) Run your forward geological model + # 2) Build the new interpolation input + interp_input = set_interp_input_fn(samples, geo_model) + + # 3) Run GemPy forward simulation simulated: gp.data.Solutions = run_gempy_forward( - interp_input=set_interp_input_fn(samples, geo_model), + interp_input=interp_input, geo_model=geo_model ) - # 4) Build likelihood and observe + # 4) Wrap in likelihood & observe lik_dist = likelihood_fn(simulated) pyro.sample(obs_name, lik_dist, obs=obs_data) return model + def make_generic_pyro_model( priors: Dict[str, Distribution], forward_fn: Callable[..., torch.Tensor], From 1cf5f0ae37f6bf1e5b693628a02b77f89f44e1bf Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 14:58:15 +0100 Subject: [PATCH 09/14] [BUG] Move backend setting to forward function for consistency Relocated `change_backend_gempy` call to `run_gempy_forward` to ensure proper backend configuration during execution. Removed redundant imports to clean up the module. --- gempy_probability/modules/forwards/_gempy_fw.py | 9 +++------ tests/test_prob_model/test_prob_factory.py | 1 - 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/gempy_probability/modules/forwards/_gempy_fw.py b/gempy_probability/modules/forwards/_gempy_fw.py index b55f7f4..3eaa602 100644 --- a/gempy_probability/modules/forwards/_gempy_fw.py +++ b/gempy_probability/modules/forwards/_gempy_fw.py @@ -1,14 +1,11 @@ -import torch -from pyro.distributions import Distribution - -import gempy_engine -from gempy.modules.data_manipulation import interpolation_input_from_structural_frame 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: - + 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 index 4781718..d18fa06 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -50,7 +50,6 @@ def test_prob_model_factory() -> None: gempy_model=geo_model, validate_serialization=False ) - BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) # 1) Define your priors model_priors = { From 2a222b15c27734f3f2db969a00653dcf4ac19c3c Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 15:11:40 +0100 Subject: [PATCH 10/14] [ENH] Introduce GemPyPyroModel type for enhanced type clarity Refactored the code to define and utilize the `GemPyPyroModel` type alias for improved type clarity across the codebase. Updated method signatures and relevant imports to align with the new type definition. --- gempy_probability/__init__.py | 2 +- .../modules/model_definition/prob_model_factory.py | 9 ++++++--- tests/test_prob_model/test_prob_factory.py | 4 ++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py index 7ffc48e..de87f4d 100644 --- a/gempy_probability/__init__.py +++ b/gempy_probability/__init__.py @@ -1,4 +1,4 @@ -from .modules.model_definition.prob_model_factory import make_gempy_pyro_model +from .modules.model_definition.prob_model_factory import make_gempy_pyro_model, GemPyPyroModel from .modules import likelihoods from ._version import __version__ diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py index ea6360d..91d5305 100644 --- a/gempy_probability/modules/model_definition/prob_model_factory.py +++ b/gempy_probability/modules/model_definition/prob_model_factory.py @@ -1,13 +1,16 @@ import pyro import torch -from typing import Callable, Dict from pyro.distributions import Distribution +from typing import Callable, Dict -from gempy_probability.modules.forwards import run_gempy_forward import gempy as gp +from gempy_probability.modules.forwards import run_gempy_forward + +GemPyPyroModel = Callable[[gp.data.GeoModel, torch.Tensor], None] def make_gempy_pyro_model( + *, priors: Dict[str, Distribution], set_interp_input_fn: Callable[ [Dict[str, torch.Tensor], gp.data.GeoModel], @@ -15,7 +18,7 @@ def make_gempy_pyro_model( ], likelihood_fn: Callable[[gp.data.Solutions], Distribution], obs_name: str = "obs" -) -> Callable[[gp.data.GeoModel, torch.Tensor], None]: +) -> GemPyPyroModel: """ Factory to produce a Pyro model for GemPy forward simulations. diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index d18fa06..79b6738 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -59,7 +59,7 @@ def test_prob_model_factory() -> None: ) } - pyro_gravity_model = gpp.make_gempy_pyro_model( + pyro_gravity_model: gpp.GemPyPyroModel = gpp.make_gempy_pyro_model( priors=model_priors, set_interp_input_fn=modify_z_for_surface_point1, likelihood_fn=gpp.likelihoods.thickness_likelihood, @@ -91,7 +91,7 @@ def modify_z_for_surface_point1( return interp_input -def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable, +def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel, y_obs_list: torch.Tensor) -> None: # Run prior sampling and visualization from pyro.infer import Predictive From 0c2a61ace206383f3e6f525c40258be5e1f18be6 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 15:51:36 +0100 Subject: [PATCH 11/14] [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. --- gempy_probability/__init__.py | 1 + .../api/model_runner/__init__.py | 1 + .../api/model_runner/_pyro_runner.py | 37 +++++++++++++++++++ gempy_probability/core/sampler.py | 0 .../model_definition/prob_model_factory.py | 6 ++- .../modules/model_runner/__init__.py | 0 tests/test_prob_model/test_prob_factory.py | 20 +++++----- 7 files changed, 53 insertions(+), 12 deletions(-) create mode 100644 gempy_probability/api/model_runner/__init__.py create mode 100644 gempy_probability/api/model_runner/_pyro_runner.py create mode 100644 gempy_probability/core/sampler.py create mode 100644 gempy_probability/modules/model_runner/__init__.py diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py index de87f4d..24eb40e 100644 --- a/gempy_probability/__init__.py +++ b/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__ diff --git a/gempy_probability/api/model_runner/__init__.py b/gempy_probability/api/model_runner/__init__.py new file mode 100644 index 0000000..97c5f4a --- /dev/null +++ b/gempy_probability/api/model_runner/__init__.py @@ -0,0 +1 @@ +from ._pyro_runner import run_predictive \ No newline at end of file diff --git a/gempy_probability/api/model_runner/_pyro_runner.py b/gempy_probability/api/model_runner/_pyro_runner.py new file mode 100644 index 0000000..de5e8a4 --- /dev/null +++ b/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 diff --git a/gempy_probability/core/sampler.py b/gempy_probability/core/sampler.py new file mode 100644 index 0000000..e69de29 diff --git a/gempy_probability/modules/model_definition/prob_model_factory.py b/gempy_probability/modules/model_definition/prob_model_factory.py index 91d5305..eb99915 100644 --- a/gempy_probability/modules/model_definition/prob_model_factory.py +++ b/gempy_probability/modules/model_definition/prob_model_factory.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] = {} diff --git a/gempy_probability/modules/model_runner/__init__.py b/gempy_probability/modules/model_runner/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index 79b6738..2f1403e 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_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())}") + + From 122e2a849e28000b178767ad9dae4225d71b0a70 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 15:57:31 +0100 Subject: [PATCH 12/14] [WIP] --- gempy_probability/core/sampler.py | 0 gempy_probability/core/samplers_data.py | 12 ++++++++++++ 2 files changed, 12 insertions(+) delete mode 100644 gempy_probability/core/sampler.py create mode 100644 gempy_probability/core/samplers_data.py diff --git a/gempy_probability/core/sampler.py b/gempy_probability/core/sampler.py deleted file mode 100644 index e69de29..0000000 diff --git a/gempy_probability/core/samplers_data.py b/gempy_probability/core/samplers_data.py new file mode 100644 index 0000000..b56e3a5 --- /dev/null +++ b/gempy_probability/core/samplers_data.py @@ -0,0 +1,12 @@ +import dataclasses + + +@dataclasses.dataclass +class NUTS_Args: + 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 \ No newline at end of file From 460bf8b000a9c81cbecd8cc24a09e5183b2e018a Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 16:38:59 +0100 Subject: [PATCH 13/14] [WIP/CLN/ENH] Extracting samplers to the api to make it easier --- gempy_probability/__init__.py | 2 +- .../api/model_runner/__init__.py | 2 +- .../api/model_runner/_pyro_runner.py | 64 +++++++++++++ gempy_probability/core/samplers_data.py | 2 +- .../modules/forwards/_gempy_fw.py | 2 +- tests/test_prob_model/test_prob_factory.py | 91 +++++++++++++------ 6 files changed, 129 insertions(+), 34 deletions(-) diff --git a/gempy_probability/__init__.py b/gempy_probability/__init__.py index 24eb40e..a695542 100644 --- 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 index 97c5f4a..c9418b3 100644 --- a/gempy_probability/api/model_runner/__init__.py +++ b/gempy_probability/api/model_runner/__init__.py @@ -1 +1 @@ -from ._pyro_runner import run_predictive \ No newline at end of file +from ._pyro_runner import run_predictive, run_nuts_inference, run_mcmc_for_NUTS \ No newline at end of file diff --git a/gempy_probability/api/model_runner/_pyro_runner.py b/gempy_probability/api/model_runner/_pyro_runner.py index de5e8a4..e65b6af 100644 --- 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 index b56e3a5..0f6178e 100644 --- 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 index 3eaa602..613a1b8 100644 --- 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 index 2f1403e..bad1dca 100644 --- 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())}") - - From 619caafe4c50f86b9eba12f9900fa6df77554346 Mon Sep 17 00:00:00 2001 From: MigueldelaVarga Date: Tue, 27 May 2025 16:45:37 +0100 Subject: [PATCH 14/14] [CLN] --- .../api/model_runner/_pyro_runner.py | 23 ++--- tests/test_prob_model/test_prob_factory.py | 97 ++++++------------- 2 files changed, 37 insertions(+), 83 deletions(-) diff --git a/gempy_probability/api/model_runner/_pyro_runner.py b/gempy_probability/api/model_runner/_pyro_runner.py index e65b6af..1133957 100644 --- a/gempy_probability/api/model_runner/_pyro_runner.py +++ b/gempy_probability/api/model_runner/_pyro_runner.py @@ -1,25 +1,14 @@ -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 +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 -from gempy_probability.core.samplers_data import NUTSConfig +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, diff --git a/tests/test_prob_model/test_prob_factory.py b/tests/test_prob_model/test_prob_factory.py index bad1dca..dd46c57 100644 --- a/tests/test_prob_model/test_prob_factory.py +++ b/tests/test_prob_model/test_prob_factory.py @@ -4,10 +4,12 @@ import pyro.distributions as dist import torch from pyro.distributions import Distribution +import arviz as az +import matplotlib.pyplot as plt import gempy as gp -from gempy_engine.core.backend_tensor import BackendTensor import gempy_probability as gpp +from gempy_engine.core.backend_tensor import BackendTensor from gempy_engine.core.data.interpolation_input import InterpolationInput from gempy_probability.core.samplers_data import NUTSConfig @@ -29,9 +31,6 @@ def test_prob_model_factory() -> None: ) ) - # TODO: Make this a more elegant way - BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH) - # TODO: Convert this into an options preset geo_model.interpolation_options.uni_degree = 0 geo_model.interpolation_options.mesh_extraction = False @@ -65,15 +64,36 @@ def test_prob_model_factory() -> None: priors=model_priors, set_interp_input_fn=modify_z_for_surface_point1, likelihood_fn=gpp.likelihoods.thickness_likelihood, - obs_name="obs_gravity" + obs_name="wells_thickness" ) - _prob_run( + data: az.InferenceData = _prob_run( geo_model=geo_model, prob_model=pyro_gravity_model, y_obs_list=torch.tensor([200, 210, 190]) ) + if False: # * Save the arviz data + data.to_netcdf("arviz_data.nc") + + _plot(data) + + # Test posterior mean values + posterior_top_mean = float(data.posterior[r'$\mu_{top}$'].mean()) + target_top = 0.00875 + target_thickness = 223 + assert abs(posterior_top_mean - target_top) < 0.0200, f"Top layer mean {posterior_top_mean} outside expected range" + posterior_thickness_mean = float(data.posterior_predictive[r'$\mu_{thickness}$'].mean()) + assert abs(posterior_thickness_mean - target_thickness) < 5, f"Thickness mean {posterior_thickness_mean} outside expected range" + # Test convergence diagnostics + assert float(data.sample_stats.diverging.sum()) == 0, "MCMC sampling has divergences" + + print("Posterior mean values:") + print(f"Top layer mean: {posterior_top_mean}") + print(f"Thickness mean: {posterior_thickness_mean}") + print("MCMC convergence diagnostics:") + print(f"Divergences: {float(data.sample_stats.diverging.sum())}") + def modify_z_for_surface_point1( samples: dict[str, Distribution], @@ -87,23 +107,15 @@ def modify_z_for_surface_point1( 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 - values=(samples[prior_key]) - ) + values=(samples[prior_key] + ) interp_input.surface_points.sp_coords = new_tensor return interp_input def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel, - y_obs_list: torch.Tensor) -> None: + y_obs_list: torch.Tensor) -> az.InferenceData: # Run prior sampling and visualization - 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 # region prior sampling prior_inference = gpp.run_predictive( @@ -133,42 +145,12 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel, plot_trace=False, run_posterior_predictive=True ) - # 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) + return data # endregion - # print("Number of interpolations: ", geo_model.counter) - - if False: # * Save the arviz data - data.to_netcdf("arviz_data.nc") +def _plot(data): az.plot_trace(data) plt.show() from gempy_probability.modules.plot.plot_posterior import default_red, default_blue @@ -188,7 +170,6 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel, colors=[default_red, default_blue], ) plt.show() - az.plot_density( data=[data.posterior_predictive, data.prior], shade=.9, @@ -197,19 +178,3 @@ def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel, 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 - target_thickness = 223 - assert abs(posterior_top_mean - target_top) < 0.0200, f"Top layer mean {posterior_top_mean} outside expected range" - posterior_thickness_mean = float(data.posterior_predictive[r'$\mu_{thickness}$'].mean()) - assert abs(posterior_thickness_mean - target_thickness) < 5, f"Thickness mean {posterior_thickness_mean} outside expected range" - # Test convergence diagnostics - assert float(data.sample_stats.diverging.sum()) == 0, "MCMC sampling has divergences" - - print("Posterior mean values:") - print(f"Top layer mean: {posterior_top_mean}") - print(f"Thickness mean: {posterior_thickness_mean}") - print("MCMC convergence diagnostics:") - print(f"Divergences: {float(data.sample_stats.diverging.sum())}")