[CLN/ENH/WIP] Figuring out extracting model definition

Author: Leguark

gempy_probability/modules/model_definition/__init__.py

@@ -0,0 +1,39 @@
+import gempy.core.data
+import gempy_engine
+from gempy.modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
+
+
+def model(geo_model: gempy.core.data.GeoModel, sp_coords_copy, 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 
+    computes the thickness of the geological layer as an observed variable.
+    """
+    # Define prior for the top layer's location:
+    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)))
+
+    # 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
+    )
+
+    # Compute the geological model
+    geo_model.solutions = gempy_engine.compute_model(
+        interpolation_input=interpolation_input,
+        options=geo_model.interpolation_options,
+        data_descriptor=geo_model.input_data_descriptor,
+        geophysics_input=geo_model.geophysics_input,
+    )
+
+    # 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)

gempy_probability/modules/model_definition/model_examples.py

@@ -44,57 +44,13 @@ def test_basic_gempy_I():
 
     from gempy_engine.core.data.interpolation_input import InterpolationInput
     from gempy_engine.core.backend_tensor import BackendTensor
-    
     from gempy.modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
 
     interpolation_input_copy: InterpolationInput = interpolation_input_from_structural_frame(geo_model)
     sp_coords_copy = interpolation_input_copy.surface_points.sp_coords
     # Change the backend to PyTorch for probabilistic modeling
     BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)
 
-    # Defining the Probabilistic Model
-    # --------------------------------
-    # The Pyro model represents the probabilistic aspects of the geological model.
-    # It defines a prior distribution for the top layer's location and computes the thickness
-    # of the geological layer as an observed variable.
-
-    def model(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 
-        computes the thickness of the geological layer as an observed variable.
-        """
-        # Define prior for the top layer's location:
-        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)))
-
-        # 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
-        )
-
-        # Compute the geological model
-        geo_model.solutions = gempy_engine.compute_model(
-            interpolation_input=interpolation_input,
-            options=geo_model.interpolation_options,
-            data_descriptor=geo_model.input_data_descriptor,
-            geophysics_input=geo_model.geophysics_input,
-        )
-
-        # 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)
-
-
-
     # %%
     # Running Prior Sampling and Visualization
     # ----------------------------------------
@@ -112,8 +68,9 @@ def model(y_obs_list):
     import pyro
     import arviz as az
     import matplotlib.pyplot as plt
-    
-    prior = Predictive(model, num_samples=50)(y_obs_list)
+
+    from gempy_probability.modules.model_definition.model_examples import model
+    prior = Predictive(model, num_samples=50)(geo_model, sp_coords_copy, y_obs_list)
 
     data = az.from_pyro(prior=prior)
     az.plot_trace(data.prior)
@@ -133,11 +90,11 @@ def model(y_obs_list):
         init_strategy=init_to_mean
     )
     mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=50, disable_validation=False)
-    mcmc.run(y_obs_list)
+    mcmc.run(geo_model, sp_coords_copy, y_obs_list)
 
 
     posterior_samples = mcmc.get_samples()
-    posterior_predictive = Predictive(model, posterior_samples)(y_obs_list)
+    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)
     az.plot_trace(data)
     plt.show()