[ENH] Introduce GemPyPyroModel type for enhanced type clarity

Refactored the code to define and utilize the `GemPyPyroModel` type alias for improved type clarity across the codebase. Updated method signatures and relevant imports to align with the new type definition.

Author: Leguark

gempy_probability/__init__.py

@@ -1,4 +1,4 @@
-from .modules.model_definition.prob_model_factory import make_gempy_pyro_model
+from .modules.model_definition.prob_model_factory import make_gempy_pyro_model, GemPyPyroModel
 from .modules import likelihoods
 
 from ._version import __version__

gempy_probability/modules/model_definition/prob_model_factory.py

@@ -1,21 +1,24 @@
 import pyro
 import torch
-from typing import Callable, Dict
 from pyro.distributions import Distribution
+from typing import Callable, Dict
 
-from gempy_probability.modules.forwards import run_gempy_forward
 import gempy as gp
+from gempy_probability.modules.forwards import run_gempy_forward
+
+GemPyPyroModel = Callable[[gp.data.GeoModel, torch.Tensor], None]
 
 
 def make_gempy_pyro_model(
+        *,
         priors: Dict[str, Distribution],
         set_interp_input_fn: Callable[
             [Dict[str, torch.Tensor], gp.data.GeoModel],
             gp.data.InterpolationInput
         ],
         likelihood_fn: Callable[[gp.data.Solutions], Distribution],
         obs_name: str = "obs"
-) -> Callable[[gp.data.GeoModel, torch.Tensor], None]:
+) -> GemPyPyroModel:
     """
     Factory to produce a Pyro model for GemPy forward simulations.
 

tests/test_prob_model/test_prob_factory.py

@@ -59,7 +59,7 @@ def test_prob_model_factory() -> None:
             )
     }
 
-    pyro_gravity_model = gpp.make_gempy_pyro_model(
+    pyro_gravity_model: gpp.GemPyPyroModel = gpp.make_gempy_pyro_model(
         priors=model_priors,
         set_interp_input_fn=modify_z_for_surface_point1,
         likelihood_fn=gpp.likelihoods.thickness_likelihood,
@@ -91,7 +91,7 @@ def modify_z_for_surface_point1(
     return interp_input
 
 
-def _prob_run(geo_model: gp.data.GeoModel, prob_model: callable,
+def _prob_run(geo_model: gp.data.GeoModel, prob_model: gpp.GemPyPyroModel,
               y_obs_list: torch.Tensor) -> None:
     # Run prior sampling and visualization
     from pyro.infer import Predictive