[CLN/ENH] Better definition for specific gempy_pyro models

Author: Leguark

gempy_probability/__init__.py

@@ -1 +1,3 @@
+from .modules.model_definition.prob_model_factory import make_gempy_pyro_model
+
 from ._version import __version__

gempy_probability/modules/forwards/__init__.py

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

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
+
+

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)

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"
-)
-

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
+    return model

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