[ENH] Refactor gravity inversion pipeline and add flexibility

Refactored `setup_geophysics` to accept dynamic density inputs and improved `calculate_scale_shift` with robust validation and detailed docstrings. Simplified the gravity inversion example by removing redundant steps and adopting modular utilities.

Author: Leguark

docs/dev_logs/2025_May.md

@@ -48,9 +48,9 @@
   - This would be useful to be able to visualize whatever the likelihood function represents independent of the probabilistic model.
   - 
 ### What would I want to have to do StonePark example as flowless as possible? 
-1. A way to save the model already with the correct **nugget effects** instead of having to "load nuggets" (TODO in gempy)
-2. Likelihood function plug and play
-3. Some sort of pre-analysis if the model looks good. Maybe some metrics about how long it takes to run the model. 
+1. [x] A way to save the model already with the correct **nugget effects** instead of having to "load nuggets" (TODO in gempy)
+2. [ ] Likelihood function plug and play
+3. [ ] Some sort of pre-analysis if the model looks good. Maybe some metrics about how long it takes to run the model. 
    1. Condition number
    2. Acceptance rate of the posterior
 4. Easy pykeops setting using .env

examples/examples/_aux_func_gravity_inversion.py

@@ -80,9 +80,11 @@ def initialize_geo_model(structural_elements: list[gp.data.StructuralElement], e
 
     return geo_model
 
-def setup_geophysics(env_path: str, geo_model: gp.data.GeoModel):
+def setup_geophysics(env_path: str, geo_model: gp.data.GeoModel, densities: list[float]):
     import pandas as pd
     from dotenv import dotenv_values
+    from gempy_engine.core.backend_tensor import BackendTensor
+    
     config = dotenv_values()
 
     df = pd.read_csv(
@@ -97,8 +99,7 @@ def setup_geophysics(env_path: str, geo_model: gp.data.GeoModel):
     interesting_columns = df[['X', 'Y', 'Bouguer_267_complete']]
     # %%
     device_location = interesting_columns[['X', 'Y']]
-    # stack 0 to the z axis
-    device_location['Z'] = 0
+    device_location['Z'] = 0 # * stack 0 to the z axis
 
     gp.set_centered_grid(
         grid=geo_model.grid,
@@ -110,10 +111,9 @@ def setup_geophysics(env_path: str, geo_model: gp.data.GeoModel):
     gravity_gradient = gp.calculate_gravity_gradient(geo_model.grid.centered_grid)
 
     # %%
-    from gempy_engine.core.backend_tensor import BackendTensor
     geo_model.geophysics_input = gp.data.GeophysicsInput(
         tz=BackendTensor.t.array(gravity_gradient),
-        densities=BackendTensor.t.array([2.61, 2.92, 3.1, 2.92, 2.61, 2.61]),
+        densities=BackendTensor.t.array(densities),
     )
 
     return interesting_columns

examples/examples/gravity_inversion.py

@@ -4,78 +4,45 @@
 
 This example demonstrates a probabilistic inversion of a geological model using Bayesian methods.
 """
-
+import numpy as np
 import os
-import time
 
 import arviz as az
-import numpy as np
+import dotenv
+import gempy_engine
+import gempy_viewer as gpv
 import pyro
 import pyro.distributions as dist
 import torch
-import xarray as xr
-from dotenv import dotenv_values
+from gempy_engine.core.backend_tensor import BackendTensor
 from matplotlib import pyplot as plt
 from pyro.infer import MCMC, NUTS, Predictive
 
 import gempy as gp
-import gempy_engine
-import gempy_viewer as gpv
-
-from examples.examples._aux_func_gravity_inversion import calculate_scale_shift, gaussian_kernel, initialize_geo_model, setup_geophysics 
-from examples.examples._aux_func_gravity_inversion_II import process_file
-
-from gempy_engine.core.backend_tensor import BackendTensor
+from examples.examples._aux_func_gravity_inversion import setup_geophysics
+from gempy_probability.modules.likelihoods.gravity_likelihoods import calculate_scale_shift, gaussian_kernel
 from gempy_probability.modules.plot.plot_posterior import default_red, default_blue
 
+dotenv.load_dotenv()
+
 # %%
 # Config
 seed = 123456
 torch.manual_seed(seed)
 pyro.set_rng_seed(seed)
 
-# %%
-# Start the timer for benchmarking purposes
-start_time = time.time()
-
-# %%
-# Load necessary configuration and paths from environment variables
-config = dotenv_values()
-path = config.get("PATH_TO_MODEL_1_Subsurface")
-
-# %%
-# Initialize lists to store structural elements for the geological model
-structural_elements = []
-global_extent = None
-color_gen = gp.data.ColorsGenerator()
-
-# %%
-# Process each .nc file in the specified directory for model construction
-for filename in os.listdir(path):
-    base, ext = os.path.splitext(filename)
-    if ext == '.nc':
-        structural_element, global_extent = process_file(os.path.join(path, filename), global_extent, color_gen)
-        structural_elements.append(structural_element)
-
-# %%
-# Configure GemPy for geological modeling with PyTorch backend
-BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH, dtype="float64")
-
-geo_model = initialize_geo_model(
-    structural_elements=structural_elements,
-    extent=(np.array(global_extent)),
-    topography=(xr.open_dataset(os.path.join(path, "Topography.nc"))),
-    load_nuggets=False
-)
-
-# TODO: Here is where we are setting the nuggets
-# * ============================================
+model_path = os.getenv("PATH_TO_NUGGET_TEST_MODEL")
+model_file = os.path.join(model_path, "nugget_effect_optimization.gempy")
+geo_model = gp.API.load_model(model_file)
 
 # %%
+# Forward Gravity
+# ---------------
 # Setup geophysics configuration for the model
 geophysics_input = setup_geophysics(
     env_path="PATH_TO_MODEL_1_BOUGUER",
-    geo_model=geo_model
+    geo_model=geo_model,
+    densities=[2.61, 2.92, 3.1, 2.92, 2.61, 2.61, 2.61]
 )
 
 # %%
@@ -87,12 +54,13 @@
 
 # %%
 # Compute the geological model
-sol = gp.compute_model(
+sol: gp.data.Solutions = gp.compute_model(
     gempy_model=geo_model,
     engine_config=gp.data.GemPyEngineConfig(
         backend=gp.data.AvailableBackends.PYTORCH,
         dtype='float64'
-    )
+    ),
+    validate_serialization=False # TODO: [ ] Validate this serialization
 )
 
 # %%
@@ -106,12 +74,11 @@
 
 # %%
 # Calculate and adapt the observed gravity data for model comparison
-grav = sol.gravity
-s, c = calculate_scale_shift(
-    a=geophysics_input["Bouguer_267_complete"].values,
-    b=(grav.detach().numpy())
+scale_factor, shift_term = calculate_scale_shift(
+    measured=geophysics_input["Bouguer_267_complete"].values,
+    simulated=sol.gravity.detach().numpy()
 )
-adapted_observed_grav = s * geophysics_input["Bouguer_267_complete"] + c
+adapted_observed_grav: np.dtype[float] = scale_factor * geophysics_input["Bouguer_267_complete"] + shift_term
 
 # %%
 # Plot the 2D gravity data for visualization
@@ -127,7 +94,10 @@
 plt.colorbar(sc, label="mGal")
 plt.show()
 
+
 # %%
+# Define Probabilistic model
+# --------------------------
 # Define hyperparameters for the Bayesian geological model
 length_scale_prior = torch.tensor(1_000.0)
 variance_prior = torch.tensor(25.0 ** 2)
@@ -141,6 +111,7 @@
     densities=prior_tensor,
 )
 
+bar
 
 # %%
 # Define the Pyro probabilistic model for inversion

gempy_probability/modules/likelihoods/gravity_likelihoods/_utils.py

@@ -1,20 +1,66 @@
 import numpy as np
-from sklearn.linear_model import LinearRegression
+from typing import Tuple
 
-def calculate_scale_shift(a: np.ndarray, b: np.ndarray) -> tuple:
-    # Reshape arrays for sklearn
-    a_reshaped = a.reshape(-1, 1)
-    b_reshaped = b.reshape(-1, 1)
+def calculate_scale_shift(
+    measured: np.ndarray,
+    simulated: np.ndarray
+) -> Tuple[float, float]:
+    """
+    Compute the best-fit scale (slope) and shift (intercept) so that:
+        simulated ≈ scale_factor * measured + shift_term
 
-    # Linear regression
-    model = LinearRegression().fit(a_reshaped, b_reshaped)
+    Parameters
+    ----------
+    measured : array-like, shape (n,)
+        The observed (independent) values.
+    simulated : array-like, shape (n,)
+        The simulated or target (dependent) values.
 
-    # Get scale and shift
-    s = model.coef_[0][0]
-    c = model.intercept_[0]
+    Returns
+    -------
+    scale_factor : float
+        Multiplicative factor.
+    shift_term : float
+        Additive offset.
 
-    return s, c
+    Raises
+    ------
+    ValueError
+        If lengths differ or fewer than 2 valid points remain.
+    """
+    # Flatten inputs
+    measured = np.asarray(measured).ravel()
+    simulated = np.asarray(simulated).ravel()
 
+    # Check same length
+    if measured.size != simulated.size:
+        raise ValueError(
+            f"Length mismatch: measured has {measured.size}, "
+            f"simulated has {simulated.size}"
+        )
+
+    # Mask out NaN or infinite entries
+    valid_mask = np.isfinite(measured) & np.isfinite(simulated)
+    measured_valid = measured[valid_mask]
+    simulated_valid = simulated[valid_mask]
+
+    if measured_valid.size < 2:
+        raise ValueError(
+            "Need at least 2 valid (non-NaN/Inf) data points"
+        )
+
+    # Build design matrix [measured_valid, 1]
+    design_matrix = np.vstack([
+        measured_valid,
+        np.ones_like(measured_valid)
+    ]).T
+
+    # Solve least squares: design_matrix @ [scale_factor, shift_term] ≈ simulated_valid
+    (scale_factor, shift_term), *_ = np.linalg.lstsq(design_matrix,
+                                                    simulated_valid,
+                                                    rcond=None)
+
+    return float(scale_factor), float(shift_term)
 
 def gaussian_kernel(locations, length_scale, variance):
     import torch
@@ -24,4 +70,3 @@ def gaussian_kernel(locations, length_scale, variance):
     # Compute the covariance matrix using the Gaussian kernel
     covariance_matrix = variance * torch.exp(-0.5 * distance_squared / length_scale ** 2)
     return covariance_matrix
-