[ENH] Add nugget optimization module with per-group optimization capability (#1036)

# Description
Refactored the nugget optimization functionality into a dedicated module for better organization and reusability. The PR extracts the optimization logic from `compute_API.py` into a new `optimize_nuggets` module with improved structure and cleaner implementation. It also adds a new API function `optimize_nuggets` that allows optimizing specific structural groups rather than the entire model.

The PR also includes code reorganization, moving the `_LoadDEMArtificial` class from `create_topography.py` to `topography.py` where it's more logically placed, and fixing import paths throughout the codebase.

Relates to #optimization-refactoring

# Checklist
- [x] My code uses type hinting for function and method arguments and return values.
- [x] I have created tests which cover my code.
- [x] The test code either 1. demonstrates at least one valuable use case (e.g. integration tests) 
or 2. verifies that outputs are as expected for given inputs (e.g. unit tests).
- [x] New tests pass locally with my changes.

Author: Leguark

gempy/API/compute_API.py

@@ -1,20 +1,19 @@
 import dotenv
+import numpy as np
 import os
-
 from typing import Optional
 
-import numpy as np
-
 import gempy_engine
-from gempy_engine.core.backend_tensor import BackendTensor
 from gempy.API.gp2_gp3_compatibility.gp3_to_gp2_input import gempy3_to_gempy2
 from gempy_engine.config import AvailableBackends
+from gempy_engine.core.backend_tensor import BackendTensor
 from gempy_engine.core.data import Solutions
-from gempy_engine.core.data.interpolation_input import InterpolationInput
 from .grid_API import set_custom_grid
+from ..core.data import StructuralGroup
 from ..core.data.gempy_engine_config import GemPyEngineConfig
 from ..core.data.geo_model import GeoModel
-from ..modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
+from ..modules.data_manipulation import interpolation_input_from_structural_frame
+from ..modules.optimize_nuggets import nugget_optimizer
 from ..optional_dependencies import require_gempy_legacy
 
 dotenv.load_dotenv()
@@ -92,91 +91,29 @@ def compute_model_at(gempy_model: GeoModel, at: np.ndarray,
     return sol.raw_arrays.custom
 
 
-def optimize_and_compute(geo_model: GeoModel, engine_config: GemPyEngineConfig, max_epochs: int = 10,
-                         convergence_criteria: float = 1e5):
+def optimize_nuggets(geo_model: GeoModel, engine_config: GemPyEngineConfig, max_epochs: int = 10,
+                     convergence_criteria: float = 1e5, only_groups:list[StructuralGroup] | None = None) -> GeoModel:
+    """
+    Optimize the nuggets of the interpolation input of the provided model.  
+    """
+
     if engine_config.backend != AvailableBackends.PYTORCH:
         raise ValueError(f'Only PyTorch backend is supported for optimization. Received {engine_config.backend}')
-
-    BackendTensor.change_backend_gempy(
-        engine_backend=engine_config.backend,
-        use_gpu=engine_config.use_gpu,
-        dtype=engine_config.dtype
-    )
-
-    import torch
-    from gempy_engine.core.data.continue_epoch import ContinueEpoch
-    interpolation_input: InterpolationInput = interpolation_input_from_structural_frame(geo_model)
-
-    geo_model.taped_interpolation_input = interpolation_input
-
-    nugget_effect_scalar: torch.Tensor = geo_model.taped_interpolation_input.surface_points.nugget_effect_scalar
-
-    optimizer = torch.optim.Adam(
-        params=[nugget_effect_scalar],
-        lr=0.01,
+    
+    geo_model = nugget_optimizer(
+        target_cond_num=convergence_criteria,
+        engine_cfg=engine_config,
+        model=geo_model,
+        max_epochs=max_epochs,
+        only_groups=only_groups
     )
 
-    # Optimization loop
-    geo_model.interpolation_options.kernel_options.optimizing_condition_number = True
-
-    def _check_convergence_criterion(conditional_number: float, condition_number_old: float, conditional_number_target: float = 1e5):
-        reached_conditional_target = conditional_number < conditional_number_target
-        if reached_conditional_target == False and epoch > 10:
-            condition_number_change = torch.abs(conditional_number - condition_number_old) / condition_number_old
-            if condition_number_change < 0.01:
-                reached_conditional_target = True
-        return reached_conditional_target
-
-    previous_condition_number = 0
-    for epoch in range(max_epochs):
-        optimizer.zero_grad()
-        try:
-            # geo_model.taped_interpolation_input.grid = geo_model.interpolation_input_copy.grid
-
-            gempy_engine.compute_model(
-                interpolation_input=geo_model.taped_interpolation_input,
-                options=geo_model.interpolation_options,
-                data_descriptor=geo_model.input_data_descriptor,
-                geophysics_input=geo_model.geophysics_input,
-            )
-        except ContinueEpoch:
-            # Get absolute values of gradients
-            grad_magnitudes = torch.abs(nugget_effect_scalar.grad)
-
-            # Get indices of the 10 largest gradients
-            grad_magnitudes.size
-
-            # * This ignores 90 percent of the gradients
-            # To int
-            n_values = int(grad_magnitudes.size()[0] * 0.9)
-            _, indices = torch.topk(grad_magnitudes, n_values, largest=False)
-
-            # Zero out gradients that are not in the top 10
-            mask = torch.ones_like(nugget_effect_scalar.grad)
-            mask[indices] = 0
-            nugget_effect_scalar.grad *= mask
-
-            # Update the vector
-            optimizer.step()
-            nugget_effect_scalar.data = nugget_effect_scalar.data.clamp_(min=1e-7)  # Replace negative values with 0
-
-            # optimizer.zero_grad()
-        # Monitor progress
-        if epoch % 1 == 0:
-            # print(f"Epoch {epoch}: Condition Number = {condition_number.item()}")
-            print(f"Epoch {epoch}")
-
-        if _check_convergence_criterion(
-                conditional_number=geo_model.interpolation_options.kernel_options.condition_number,
-                condition_number_old=previous_condition_number,
-                conditional_number_target=convergence_criteria,
-        ):
-            break
-        previous_condition_number = geo_model.interpolation_options.kernel_options.condition_number
-        continue
-
-    geo_model.interpolation_options.kernel_options.optimizing_condition_number = False
+    return geo_model
 
+def optimize_and_compute(geo_model: GeoModel, engine_config: GemPyEngineConfig, max_epochs: int = 10,
+                         convergence_criteria: float = 1e5):
+    
+    optimize_nuggets(geo_model, engine_config, max_epochs, convergence_criteria)
     geo_model.solutions = gempy_engine.compute_model(
         interpolation_input=geo_model.taped_interpolation_input,
         options=geo_model.interpolation_options,

gempy/core/data/geo_model.py

@@ -15,14 +15,13 @@
 from gempy_engine.core.data.interpolation_input import InterpolationInput
 from gempy_engine.core.data.raw_arrays_solution import RawArraysSolution
 from gempy_engine.core.data.transforms import Transform, GlobalAnisotropy
-from gempy_engine.modules.geophysics.gravity_gradient import calculate_gravity_gradient
 from .encoders.converters import instantiate_if_necessary
 from .encoders.json_geomodel_encoder import encode_numpy_array
 from .grid import Grid
 from .orientations import OrientationsTable
 from .structural_frame import StructuralFrame
 from .surface_points import SurfacePointsTable
-from ...modules.data_manipulation.engine_factory import interpolation_input_from_structural_frame
+from ...modules.data_manipulation import interpolation_input_from_structural_frame
 
 
 """
@@ -319,6 +318,7 @@ def deserialize_properties(cls, data: Union["GeoModel", dict], constructor: Mode
                 # * Reset geophysics if necessary
                 centered_grid = instance.grid.centered_grid
                 if centered_grid is not None and instance.geophysics_input is not None:
+                    from gempy_engine.modules.geophysics.gravity_gradient import calculate_gravity_gradient
                     instance.geophysics_input.tz = calculate_gravity_gradient(centered_grid)
 
                 return instance

gempy/core/data/grid_modules/topography.py

@@ -7,9 +7,8 @@
 import numpy as np
 
 from .regular_grid import RegularGrid
-from ....modules.grids.create_topography import _LoadDEMArtificial
 
-from ....optional_dependencies import require_skimage
+from ....optional_dependencies import require_skimage, require_scipy
 from dataclasses import field, dataclass
 from ..encoders.converters import short_array_type
 
@@ -256,3 +255,107 @@ def load(self, path):
 
     def load_from_saved(self, *args, **kwargs):
         self.load(*args, **kwargs)
+
+
+class _LoadDEMArtificial:  # * Cannot think of a good reason to be a class
+
+    def __init__(self, grid=None, fd=2.0, extent=None, resolution=None, d_z=None):
+        """Class to create a random topography based on a fractal grid algorithm.
+
+        Args:
+            fd:         fractal dimension, defaults to 2.0
+            d_z:        maximum height difference. If none, last 20% of the model in z direction
+            extent:     extent in xy direction. If none, geo_model.grid.extent
+            resolution: desired resolution of the topography array. If none, geo_model.grid.resolution
+        """
+        self.values_2d = np.array([])
+        self.resolution = grid.resolution[:2] if resolution is None else resolution
+
+        assert all(np.asarray(self.resolution) >= 2), 'The regular grid needs to be at least of size 2 on all directions.'
+        self.extent = grid.extent if extent is None else extent
+
+        if d_z is None:
+            self.d_z = np.array([self.extent[5] - (self.extent[5] - self.extent[4]) * 1 / 5, self.extent[5]])
+            print(self.d_z)
+        else:
+            self.d_z = d_z
+
+        topo = self.fractalGrid(fd, n=self.resolution.max())
+        topo = np.interp(topo, (topo.min(), topo.max()), self.d_z)
+
+        self.dem_zval = topo[:self.resolution[0], :self.resolution[1]]  # crop fractal grid with resolution
+        self.create_topo_array()
+
+    @staticmethod
+    def fractalGrid(fd, n=256):
+        """
+        Modified after https://github.com/samthiele/pycompass/blob/master/examples/3_Synthetic%20Examples.ipynb
+
+        Generate isotropic fractal surface image using
+        spectral synthesis method [1, p.]
+        References:
+        1. Yuval Fisher, Michael McGuire,
+        The Science of Fractal Images, 1988
+
+        (cf. http://shortrecipes.blogspot.com.au/2008/11/python-isotropic-fractal-surface.html)
+        **Arguments**:
+         -fd = the fractal dimension
+         -N = the size of the fractal surface/image
+
+        """
+        h = 1 - (fd - 2)
+        # X = np.zeros((N, N), complex)
+        a = np.zeros((n, n), complex)
+        powerr = -(h + 1.0) / 2.0
+
+        for i in range(int(n / 2) + 1):
+            for j in range(int(n / 2) + 1):
+                phase = 2 * np.pi * np.random.rand()
+
+                if i != 0 or j != 0:
+                    rad = (i * i + j * j) ** powerr * np.random.normal()
+                else:
+                    rad = 0.0
+
+                a[i, j] = complex(rad * np.cos(phase), rad * np.sin(phase))
+
+                if i == 0:
+                    i0 = 0
+                else:
+                    i0 = n - i
+
+                if j == 0:
+                    j0 = 0
+                else:
+                    j0 = n - j
+
+                a[i0, j0] = complex(rad * np.cos(phase), -rad * np.sin(phase))
+
+                a.imag[int(n / 2)][0] = 0.0
+                a.imag[0, int(n / 2)] = 0.0
+                a.imag[int(n / 2)][int(n / 2)] = 0.0
+
+        for i in range(1, int(n / 2)):
+            for j in range(1, int(n / 2)):
+                phase = 2 * np.pi * np.random.rand()
+                rad = (i * i + j * j) ** powerr * np.random.normal()
+                a[i, n - j] = complex(rad * np.cos(phase), rad * np.sin(phase))
+                a[n - i, j] = complex(rad * np.cos(phase), -rad * np.sin(phase))
+
+        scipy = require_scipy()
+        itemp = scipy.fftpack.ifft2(a)
+        itemp = itemp - itemp.min()
+
+        return itemp.real / itemp.real.max()
+
+    def create_topo_array(self):
+        """for masking the lith block"""
+        x = np.linspace(self.extent[0], self.extent[1], self.resolution[0])
+        y = np.linspace(self.extent[2], self.extent[3], self.resolution[1])
+        self.x = x
+        self.y = y
+        xx, yy = np.meshgrid(x, y, indexing='ij')
+        self.values_2d = np.dstack([xx, yy, self.dem_zval])
+
+    def get_values(self):
+        return self.values_2d

gempy/modules/data_manipulation/__init__.py

@@ -0,0 +1 @@
+from ._engine_factory import interpolation_input_from_structural_frame

gempy/modules/data_manipulation/engine_factory.py renamed to gempy/modules/data_manipulation/_engine_factory.py

@@ -2,9 +2,9 @@
 
 import numpy as np
 
-from gempy.core.data import GeoModel, StructuralFrame, SurfacePointsTable, StructuralElement, OrientationsTable
-from gempy.core.data.orientations import DEFAULT_ORI_NUGGET
-from gempy.core.data.surface_points import DEFAULT_SP_NUGGET
+from ...core.data import GeoModel, StructuralFrame, SurfacePointsTable, StructuralElement, OrientationsTable
+from ...core.data.orientations import DEFAULT_ORI_NUGGET
+from ...core.data.surface_points import DEFAULT_SP_NUGGET
 
 
 def add_surface_points(