[ENH] Improve CG solver and add weights caching for better performance (#17)

# Enhanced Conjugate Gradient Solver with Improved Stability and GPU Support

This PR significantly improves the robustness and performance of our solver system with several key enhancements:

## Solver Improvements
- Completely refactored the `ConjugateGradientSolver` with enhanced stability features for ill-conditioned matrices
- Added adaptive tolerance, regularization support, and robust convergence criteria
- Implemented residual monitoring and restart capability for better convergence
- Added detailed diagnostics and warnings for numerical issues

## GPU Acceleration
- Added GPU support for PyKeOps solver in the torch backend
- Properly configured backend tensor handling for GPU operations
- Fixed device selection logic in solver interface

## Caching and Weights Management
- Added ability to pass initial weights through the interpolation input
- Implemented `weights_x0` in solver input for warm starts
- Added `clear_cache` method to `WeightCache` for better memory management

## Numerical Stability
- Fixed condition number computation and made plotting optional
- Properly store condition number in kernel options during operations
- Added safeguards against numerical breakdown in denominator calculations
- Implemented stagnation detection for ill-conditioned systems

## Experimental Features
- Added Nyström preconditioner implementation (currently disabled by default)
- Included configurable pivoting strategies for matrix approximation

These changes significantly improve the solver's ability to handle challenging numerical problems while maintaining performance.

Author: Leguark

gempy_engine/API/interp_single/_interp_scalar_field.py

@@ -47,6 +47,7 @@ def interpolate_scalar_field(solver_input: SolverInput, options: InterpolationOp
 
     match weights_cached:
         case None:
+            foo = solver_input.weights_x0
             weights = _solve_and_store_weights(
                 solver_input=solver_input,
                 kernel_options=options.kernel_options,
@@ -87,7 +88,6 @@ def _solve_interpolation(interp_input: SolverInput, kernel_options: KernelOption
     if kernel_options.optimizing_condition_number:
         _optimize_nuggets_against_condition_number(A_matrix, interp_input, kernel_options)
 
-    # TODO: Smooth should be taken from options
     weights = solver_interface.kernel_reduction(
         cov=A_matrix,
         b=b_vector,

gempy_engine/API/interp_single/_interp_single_feature.py

@@ -108,6 +108,7 @@ def input_preprocess(data_shape: TensorsStructure, interpolation_input: Interpol
         xyz_to_interpolate=grid_internal,
         fault_internal=fault_values
     )
+    solver_input.weights_x0 = interpolation_input.weights
 
     return solver_input
 

gempy_engine/core/backend_tensor.py

@@ -12,6 +12,8 @@
 
 if is_pytorch_installed:
     import torch
+    
+PYKEOPS= DEFAULT_PYKEOPS
 
 # * Import a copy of numpy as tfnp
 from importlib.util import find_spec, module_from_spec
@@ -44,7 +46,7 @@ def get_backend_string(cls) -> str:
 
     @classmethod
     def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = True, dtype: Optional[str] = None):
-        cls._change_backend(engine_backend, pykeops_enabled=DEFAULT_PYKEOPS, use_gpu=use_gpu, dtype=dtype)
+        cls._change_backend(engine_backend, pykeops_enabled=PYKEOPS, use_gpu=use_gpu, dtype=dtype)
 
     @classmethod
     def _change_backend(cls, engine_backend: AvailableBackends, pykeops_enabled: bool = False, use_gpu: bool = True, dtype: Optional[str] = None):
@@ -100,6 +102,12 @@ def _change_backend(cls, engine_backend: AvailableBackends, pykeops_enabled: boo
                 if (pykeops_enabled):
                     import pykeops
                     cls._wrap_pykeops_functions()
+                
+                if (use_gpu):
+                    cls.use_gpu = True
+                    # cls.tensor_backend_pointer['active_backend'].set_default_device("cuda")
+                else:
+                    cls.use_gpu = False
 
             case (_):
                 raise AttributeError(

gempy_engine/core/data/interpolation_input.py

@@ -1,6 +1,6 @@
 import pprint
-from dataclasses import dataclass
-from typing import Optional
+from dataclasses import dataclass, field
+from typing import Optional, Union
 
 import numpy as np
 
@@ -21,23 +21,9 @@ class InterpolationInput:
     surface_points: SurfacePoints
     orientations: Orientations
     _original_grid: EngineGrid
-    
-    @property
-    def original_grid(self):
-        return self._original_grid
-    
-    def set_grid_to_original(self):
-        self._grid = self._original_grid
-
-
     _grid: EngineGrid
-    @property
-    def grid(self):
-        return self._grid
-    
-    def set_temp_grid(self, value):
-        self._grid = value
 
+    weights: Union[list[np.ndarray] | np.ndarray] = field(default_factory=lambda: [])
     _unit_values: Optional[np.ndarray] = None
     segmentation_function: Optional[callable] = None  # * From scalar field to values
 
@@ -52,14 +38,18 @@ def set_temp_grid(self, value):
 
     def __init__(self, surface_points: SurfacePoints, orientations: Orientations, grid: EngineGrid,
                  unit_values: Optional[np.ndarray] = None, segmentation_function: Optional[callable] = None,
-                 stack_relation: StackRelationType = StackRelationType.ERODE):
+                 stack_relation: StackRelationType = StackRelationType.ERODE, weights: list[np.ndarray] = None):
+        if weights is None:
+            weights = []
+        
         self.surface_points = surface_points
         self._original_grid = grid
         self._grid = grid
         self.orientations = orientations
         self.unit_values = unit_values
         self.segmentation_function = segmentation_function
         self.stack_relation = stack_relation
+        self.weights = weights
 
     # @ on
 
@@ -92,6 +82,7 @@ def from_interpolation_input_subset(cls, all_interpolation_input: "Interpolation
             grid=grid,
             unit_values=unit_values,
             stack_relation=stack_structure.active_masking_descriptor,
+            weights=(all_interpolation_input.weights[stack_number] if stack_number < len(all_interpolation_input.weights) else None)
         )
 
         # ! Setting this on the constructor does not work with data classes.
@@ -116,6 +107,19 @@ def from_schema(cls, schema: InterpolationInputSchema) -> "InterpolationInput":
             grid=grid
         )
 
+    @property
+    def original_grid(self):
+        return self._original_grid
+
+    def set_grid_to_original(self):
+        self._grid = self._original_grid
+
+    @property
+    def grid(self):
+        return self._grid
+
+    def set_temp_grid(self, value):
+        self._grid = value
 
     @property
     def slice_feature(self):