[CLN/ENH/WIP] Breaking down the different pieces of a prob model

Author: Leguark

gempy_probability/modules/model_definition/model_examples.py

@@ -2,6 +2,10 @@
 import gempy_engine
 from gempy.modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
 
+import pyro
+import pyro.distributions as dist
+import torch
+
 
 def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
     """
@@ -10,21 +14,34 @@ def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
     computes the thickness of the geological layer as an observed variable.
     """
     # Define prior for the top layer's location:
+    # region Prior definition
     prior_mean = sp_coords_copy[0, 2]
-    import pyro
-    import pyro.distributions as dist
-    import torch
-    mu_top = pyro.sample(r'$\mu_{top}$', dist.Normal(prior_mean, torch.tensor(0.02, dtype=torch.float64)))
+    normal = dist.Normal(
+        loc=prior_mean,
+        scale=torch.tensor(0.02, dtype=torch.float64)
+    )
+
+    mu_top = pyro.sample(
+        name=r'$\mu_{top}$',
+        fn=normal
+    )
+    # endregion
+
+    # region Prior injection into the gempy model
 
-    # Update the model with the new top layer's location
+    # * Update the model with the new top layer's location
     interpolation_input = interpolation_input_from_structural_frame(geo_model)
     interpolation_input.surface_points.sp_coords = torch.index_put(
-        interpolation_input.surface_points.sp_coords,
-        (torch.tensor([0]), torch.tensor([2])),
-        mu_top
+        input=interpolation_input.surface_points.sp_coords,
+        indices=(torch.tensor([0]), torch.tensor([2])),
+        values=mu_top
     )
 
-    # Compute the geological model
+    # endregion
+
+    # region Forward model computation
+
+    # * Compute the geological model
     geo_model.solutions = gempy_engine.compute_model(
         interpolation_input=interpolation_input,
         options=geo_model.interpolation_options,
@@ -35,5 +52,15 @@ def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, y_obs_list):
     # Compute and observe the thickness of the geological layer
     simulated_well = geo_model.solutions.octrees_output[0].last_output_center.custom_grid_values
     thickness = simulated_well.sum()
-    pyro.deterministic(r'$\mu_{thickness}$', thickness.detach())
-    y_thickness = pyro.sample(r'$y_{thickness}$', dist.Normal(thickness, 50), obs=y_obs_list)
+    pyro.deterministic(
+        name=r'$\mu_{thickness}$',
+        value=thickness.detach()
+    )
+
+    # endregion
+
+    y_thickness = pyro.sample(
+        name=r'$y_{thickness}$',
+        fn=dist.Normal(thickness, 50),
+        obs=y_obs_list
+    )

tests/test_prob_model/test_prob_I.py

@@ -3,6 +3,7 @@
 import gempy_engine
 import numpy as np
 
+
 def test_basic_gempy_I():
     current_dir = os.path.dirname(os.path.abspath(__file__))
     data_path = os.path.abspath(os.path.join(current_dir, '..', '..', 'examples', 'tutorials', 'data'))
@@ -21,7 +22,6 @@ def test_basic_gempy_I():
     geo_model.interpolation_options.mesh_extraction = False
     geo_model.interpolation_options.sigmoid_slope = 1100.
 
-
     x_loc = 6000
     y_loc = 0
     z_loc = np.linspace(0, 4000, 100)
@@ -70,13 +70,17 @@ def test_basic_gempy_I():
     import matplotlib.pyplot as plt
 
     from gempy_probability.modules.model_definition.model_examples import model
-    prior = Predictive(model, num_samples=50)(geo_model, sp_coords_copy, y_obs_list)
-    
+    predictive = Predictive(
+        model=model,
+        num_samples=50
+    )
+
+    prior = predictive(geo_model, sp_coords_copy, y_obs_list)
+
     data = az.from_pyro(prior=prior)
     az.plot_trace(data.prior)
     plt.show()
 
-
     from pyro.infer import NUTS
     from pyro.infer import MCMC
     from pyro.infer.autoguide import init_to_mean
@@ -92,7 +96,6 @@ def test_basic_gempy_I():
     mcmc = MCMC(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)
     data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)