[CLN] Define better likelihood module

Author: Leguark

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__

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
 

gempy_probability/modules/likelihoods/_gravity_inv.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):

gempy_probability/modules/likelihoods/_likelihood_functions.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 

gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.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():

gempy_probability/modules/model_definition/model_examples.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