[WIP] Split the logic for gravity inversion

Author: Leguark

examples/examples/_aux_func_gravity_inversion.py

@@ -1,5 +1,3 @@
-import numpy as np
-from sklearn.linear_model import LinearRegression
 
 import numpy as np
 
@@ -8,29 +6,6 @@
 from vector_geology.model_building_functions import optimize_nuggets_for_group
 
 
-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)
-
-    # Linear regression
-    model = LinearRegression().fit(a_reshaped, b_reshaped)
-
-    # Get scale and shift
-    s = model.coef_[0][0]
-    c = model.intercept_[0]
-
-    return s, c
-
-
-def gaussian_kernel(locations, length_scale, variance):
-    import torch
-    # Compute the squared Euclidean distance between each pair of points
-    locations = torch.tensor(locations.values)
-    distance_squared = torch.cdist(locations, locations, p=2).pow(2)
-    # Compute the covariance matrix using the Gaussian kernel
-    covariance_matrix = variance * torch.exp(-0.5 * distance_squared / length_scale ** 2)
-    return covariance_matrix
 
 
 def initialize_geo_model(structural_elements: list[gp.data.StructuralElement], extent: list[float],
@@ -105,56 +80,6 @@ def initialize_geo_model(structural_elements: list[gp.data.StructuralElement], e
 
     return geo_model
 
-
-def optimize_nuggets_for_whole_project(geo_model: gp.data.GeoModel):
-    geo_model.interpolation_options.kernel_options.range = 0.7
-    geo_model.interpolation_options.kernel_options.c_o = 4
-    optimize_nuggets_for_group(
-        geo_model=geo_model,
-        structural_group=geo_model.structural_frame.get_group_by_name('Red'),
-        plot_evaluation=False,
-        plot_result=True
-    )
-    geo_model.interpolation_options.kernel_options.range = 2
-    geo_model.interpolation_options.kernel_options.c_o = 4
-    optimize_nuggets_for_group(
-        geo_model=geo_model,
-        structural_group=geo_model.structural_frame.get_group_by_name('Blue'),
-        plot_evaluation=False,
-        plot_result=False
-    )
-    if False:
-        optimize_nuggets_for_group(
-            geo_model=geo_model,
-            structural_group=geo_model.structural_frame.get_group_by_name('Green'),
-            plot_evaluation=False,
-            plot_result=True
-        )
-
-
-def apply_optimized_nuggets(geo_model: gp.data.GeoModel, loaded_nuggets_red, loaded_nuggets_blue, loaded_nuggets_green):
-    gp.modify_surface_points(
-        geo_model,
-        slice=None,
-        elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Red').elements],
-        nugget=loaded_nuggets_red
-    )
-    if True:  # Ignore OB
-        gp.modify_surface_points(
-            geo_model,
-            slice=None,
-            elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Blue').elements],
-            nugget=loaded_nuggets_blue
-        )
-    if False:
-        gp.modify_surface_points(
-            geo_model,
-            slice=None,
-            elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Green').elements],
-            nugget=loaded_nuggets_green
-        )
-
-
 def setup_geophysics(env_path: str, geo_model: gp.data.GeoModel):
     import pandas as pd
     from dotenv import dotenv_values

examples/examples/_nugget_stuff.py

@@ -0,0 +1,115 @@
+import gempy as gp
+from vector_geology.model_building_functions import optimize_nuggets_for_group
+import gempy as gp
+import numpy as np
+
+
+def optimize_nuggets_for_group(geo_model: gp.data.GeoModel, structural_group: gp.data.StructuralGroup,
+                               plot_evaluation: bool = False, plot_result: bool = False) -> None:
+    temp_structural_frame = gp.data.StructuralFrame(
+        structural_groups=[structural_group],
+        color_gen=gp.data.ColorsGenerator()
+    )
+
+    previous_structural_frame = geo_model.structural_frame
+
+    geo_model.structural_frame = temp_structural_frame
+
+    gp.API.compute_API.optimize_and_compute(
+        geo_model=geo_model,
+        engine_config=gp.data.GemPyEngineConfig(
+            backend=gp.data.AvailableBackends.PYTORCH,
+        ),
+        max_epochs=100,
+        convergence_criteria=1e5
+    )
+
+    nugget_effect = geo_model.taped_interpolation_input.surface_points.nugget_effect_scalar.detach().numpy()
+    np.save(f"temp/nuggets_{structural_group.name}", nugget_effect)
+
+    if plot_evaluation:
+        import matplotlib.pyplot as plt
+
+        plt.hist(nugget_effect, bins=50, color='black', alpha=0.7, log=True)
+        plt.xlabel('Eigenvalue')
+        plt.ylabel('Frequency')
+        plt.title('Histogram of Eigenvalues (nugget-grad)')
+        plt.show()
+
+    if plot_result:
+        import gempy_viewer as gpv
+        import pyvista as pv
+
+        gempy_vista = gpv.plot_3d(
+            model=geo_model,
+            show=False,
+            kwargs_plot_structured_grid={'opacity': 0.3}
+        )
+
+        # Create a point cloud mesh
+        surface_points_xyz = geo_model.surface_points_copy.df[['X', 'Y', 'Z']].to_numpy()
+
+        point_cloud = pv.PolyData(surface_points_xyz[0:])
+        point_cloud['values'] = nugget_effect
+
+        gempy_vista.p.add_mesh(
+            point_cloud,
+            scalars='values',
+            cmap='inferno',
+            point_size=25,
+        )
+
+        gempy_vista.p.show()
+
+    geo_model.structural_frame = previous_structural_frame
+    return nugget_effect
+
+
+def optimize_nuggets_for_whole_project(geo_model: gp.data.GeoModel):
+    geo_model.interpolation_options.kernel_options.range = 0.7
+    geo_model.interpolation_options.kernel_options.c_o = 4
+    optimize_nuggets_for_group(
+        geo_model=geo_model,
+        structural_group=geo_model.structural_frame.get_group_by_name('Red'),
+        plot_evaluation=False,
+        plot_result=True
+    )
+    geo_model.interpolation_options.kernel_options.range = 2
+    geo_model.interpolation_options.kernel_options.c_o = 4
+    optimize_nuggets_for_group(
+        geo_model=geo_model,
+        structural_group=geo_model.structural_frame.get_group_by_name('Blue'),
+        plot_evaluation=False,
+        plot_result=False
+    )
+    if False:
+        optimize_nuggets_for_group(
+            geo_model=geo_model,
+            structural_group=geo_model.structural_frame.get_group_by_name('Green'),
+            plot_evaluation=False,
+            plot_result=True
+        )
+
+
+def apply_optimized_nuggets(geo_model: gp.data.GeoModel, loaded_nuggets_red, loaded_nuggets_blue, loaded_nuggets_green):
+    gp.modify_surface_points(
+        geo_model,
+        slice=None,
+        elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Red').elements],
+        nugget=loaded_nuggets_red
+    )
+    if True:  # Ignore OB
+        gp.modify_surface_points(
+            geo_model,
+            slice=None,
+            elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Blue').elements],
+            nugget=loaded_nuggets_blue
+        )
+    if False:
+        gp.modify_surface_points(
+            geo_model,
+            slice=None,
+            elements_names=[element.name for element in geo_model.structural_frame.get_group_by_name('Green').elements],
+            nugget=loaded_nuggets_green
+        )
+

examples/examples/gravity_inversion.py

@@ -60,13 +60,17 @@
 # %%
 # 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
+# * ============================================
+
 # %%
 # Setup geophysics configuration for the model
 geophysics_input = setup_geophysics(

examples/examples/gravity_inversion_I.py

@@ -0,0 +1,71 @@
+"""
+Probabilistic Inversion Example: Geological Model: Making the model stable
+--------------------------------------------------------------------------
+
+
+"""
+
+
+import os
+import time
+
+import arviz as az
+import numpy as np
+import pyro
+import pyro.distributions as dist
+import torch
+import xarray as xr
+from dotenv import dotenv_values
+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 gempy_probability.modules.plot.plot_posterior import default_red, default_blue
+
+# %%
+# 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
+)
+

gempy_probability/modules/likelihoods/gravity_likelihoods/__init__.py

@@ -0,0 +1 @@
+from ._utils import calculate_scale_shift, gaussian_kernel

gempy_probability/modules/likelihoods/gravity_likelihoods/_utils.py

@@ -0,0 +1,27 @@
+import numpy as np
+from sklearn.linear_model import LinearRegression
+
+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)
+
+    # Linear regression
+    model = LinearRegression().fit(a_reshaped, b_reshaped)
+
+    # Get scale and shift
+    s = model.coef_[0][0]
+    c = model.intercept_[0]
+
+    return s, c
+
+
+def gaussian_kernel(locations, length_scale, variance):
+    import torch
+    # Compute the squared Euclidean distance between each pair of points
+    locations = torch.tensor(locations.values)
+    distance_squared = torch.cdist(locations, locations, p=2).pow(2)
+    # Compute the covariance matrix using the Gaussian kernel
+    covariance_matrix = variance * torch.exp(-0.5 * distance_squared / length_scale ** 2)
+    return covariance_matrix
+