[WIP] A bit of model analysis

Author: Leguark

docs/dev_logs/2025_May.md

@@ -2,9 +2,12 @@
 
 --- General
 - [ ] Create a proper suit of tests
+  - [ ] Keep extracting code into modules and functions 
+  - [ ] Posterior analysis for gempy 
 - [ ] Investigate speed
 - [ ] Investigate segmentation function 
 - [ ] Create a robust api to make sure we do not mess up assigning priors to gempy parameters
+- [ ] do we need to use index_put or just normal indexing?
 --- Vector examples
 - [ ] Bring the Vector example here...
 - [ ] Using categories as likelihood function (maybe I need to check out with Sam)

gempy_probability/modules/model_definition/model_examples.py

@@ -18,7 +18,7 @@ def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
     prior_mean = sp_coords_copy[0, 2]
     normal = dist.Normal(
         loc=prior_mean,
-        scale=torch.tensor(0.02, dtype=torch.float64)
+        scale=torch.tensor(0.1, dtype=torch.float64)
     )
 
     mu_top = pyro.sample(
@@ -36,6 +36,7 @@ def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
         indices=(torch.tensor([0]), torch.tensor([2])),
         values=mu_top
     )
+    # interpolation_input.surface_points.sp_coords[0, 2] = mu_top
 
     # endregion
 
@@ -61,6 +62,6 @@ def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
 
     y_thickness = pyro.sample(
         name=r'$y_{thickness}$',
-        fn=dist.Normal(thickness, 50),
+        fn=dist.Normal(thickness, 25),
         obs=y_obs_list
     )

tests/test_prob_model/test_prob_I.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 
-def test_basic_gempy_I():
+def test_basic_gempy_I() -> 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(
@@ -84,6 +84,7 @@ def test_basic_gempy_I():
     from pyro.infer import NUTS
     from pyro.infer import MCMC
     from pyro.infer.autoguide import init_to_mean
+
     pyro.primitives.enable_validation(is_validate=True)
     nuts_kernel = NUTS(
         model,
@@ -93,11 +94,42 @@ def test_basic_gempy_I():
         max_tree_depth=10,
         init_strategy=init_to_mean
     )
-    mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=50, disable_validation=False)
+    mcmc = MCMC(
+        kernel=nuts_kernel,
+        num_samples=200,
+        warmup_steps=50,
+        disable_validation=False
+    )
     mcmc.run(geo_model, sp_coords_copy, y_obs_list)
 
     posterior_samples = mcmc.get_samples()
-    posterior_predictive = Predictive(model, posterior_samples)(geo_model, sp_coords_copy, y_obs_list)
+    
+    posterior_predictive_fn = Predictive(
+        model=model,
+        posterior_samples=posterior_samples
+    )
+
+    posterior_predictive = posterior_predictive_fn(geo_model, sp_coords_copy, y_obs_list)
+
     data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)
     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()