[ENH] Add soft segment activation function for improved layer segmentation (#11)

# Improved Activation Functions with Soft Segmentation

This PR introduces a new soft segmentation approach for activation functions with the following changes:

- Added `_soft_segment.py` module with `soft_segment_unbounded` function that provides improved sigmoid-based segmentation
- Implemented specialized segmentation functions for lithology and fault handling
- Updated the activation interface to use the new soft segmentation by default
- Added support for both scalar and array-based sigmoid slope parameters
- Incorporated per-edge adaptive temperature calculations for better control of sigmoid transitions
- Added deprecation warning to the legacy hard sigmoid activation function
- Created comprehensive test cases to validate the new segmentation functions

The new implementation provides more precise control over the transition between geological formations and supports both scalar field values and fault handling with improved numerical stability.

Author: Leguark

gempy_engine/modules/activator/_soft_segment.py

@@ -0,0 +1,93 @@
+import numbers
+
+import numpy as np
+
+from ...core.backend_tensor import BackendTensor as bt, BackendTensor
+
+try:
+    import torch
+except ModuleNotFoundError:
+    pass
+
+
+def soft_segment_unbounded(Z, edges, ids, sigmoid_slope):
+    """
+    Z:            array of shape (...,) of scalar values
+    edges:        array of shape (K-1,) of finite split points [e1, e2, ..., e_{K-1}]
+    ids:          array of shape (K,)   of the id for each of the K bins
+    sigmoid_slope: scalar target peak slope m > 0
+    returns:      array of shape (...,) of the soft-assigned id
+    """
+    ids = bt.t.array(ids[::-1].copy())
+
+    # Check if sigmoid function is num or array
+    match sigmoid_slope:
+        case numbers.Number():
+            membership = _lith_segmentation(Z, edges, ids, sigmoid_slope)
+        case _ if isinstance(sigmoid_slope, (np.ndarray, torch.Tensor)):
+            membership = _final_faults_segmentation(Z, edges, sigmoid_slope)
+        case _:
+            raise ValueError("sigmoid_slope must be a float or an array")
+
+    ids__sum = bt.t.sum(membership * ids, axis=-1)
+    return ids__sum[None, :]
+
+
+def _final_faults_segmentation(Z, edges, sigmoid_slope):
+    first = _sigmoid(
+        scalar_field=Z,
+        edges=edges[0],
+        tau_k=1 / sigmoid_slope
+    )  # shape (...,)
+    last = _sigmoid(
+        scalar_field=Z,
+        edges=edges[-1],
+        tau_k=1 / sigmoid_slope
+    )
+    membership = bt.t.concatenate(
+        [first[..., None], last[..., None]],
+        axis=-1
+    )  # shape (...,K)
+    return membership
+
+
+def _lith_segmentation(Z, edges, ids, sigmoid_slope):
+    # 1) per-edge temperatures τ_k = |Δ_k|/(4·m)
+    jumps = bt.t.abs(ids[1:] - ids[:-1])  # shape (K-1,)
+    tau_k = jumps / float(sigmoid_slope)  # shape (K-1,)
+    # 2) first bin (-∞, e1) via σ((e1 - Z)/τ₁)
+    first = _sigmoid(
+        scalar_field=-Z,
+        edges=-edges[0],
+        tau_k=tau_k[0]
+    )  # shape (...,)
+    # 3) last  bin [e_{K-1}, ∞) via σ((Z - e_{K-1})/τ_{K-1})
+    # last = 1.0 / (1.0 + np.exp(-(Z - edges[-1]) / tau_k[-1]))  # shape (...,)
+    last = _sigmoid(
+        scalar_field=Z,
+        edges=edges[-1],
+        tau_k=tau_k[-1]
+    )
+    # 4) middle bins [e_i, e_{i+1}): σ((Z - e_i)/τ_i) - σ((Z - e_{i+1})/τ_{i+1})
+    # shape (...,1)
+    left = _sigmoid(
+        scalar_field=(Z[..., None]),
+        edges=edges[:-1],
+        tau_k=tau_k[:-1]
+    )
+    right = _sigmoid(
+        scalar_field=(Z[..., None]),
+        edges=edges[1:],
+        tau_k=tau_k[1:]
+    )
+    middle = left - right  # (...,K-2)
+    # 5) assemble memberships and weight by ids
+    membership = bt.t.concatenate(
+        [first[..., None], middle, last[..., None]],
+        axis=-1
+    )  # shape (...,K)
+    return membership
+
+
+def _sigmoid(scalar_field, edges, tau_k):
+    return 1.0 / (1.0 + bt.t.exp(-(scalar_field - edges) / tau_k))

gempy_engine/modules/activator/activator_interface.py

@@ -1,25 +1,34 @@
 import warnings
 
-from gempy_engine.config import DEBUG_MODE, AvailableBackends
-from gempy_engine.core.backend_tensor import BackendTensor as bt, BackendTensor
-import numpy as np
+from ...config import DEBUG_MODE, AvailableBackends
+from ...core.backend_tensor import BackendTensor as bt, BackendTensor
+from ...core.data.exported_fields import ExportedFields
+from ._soft_segment import soft_segment_unbounded
 
-from gempy_engine.core.data.exported_fields import ExportedFields
+import numpy as np
 
 
 def activate_formation_block(exported_fields: ExportedFields, ids: np.ndarray,
                              sigmoid_slope: float) -> np.ndarray:
     Z_x: np.ndarray = exported_fields.scalar_field_everywhere
     scalar_value_at_sp: np.ndarray = exported_fields.scalar_field_at_surface_points
 
-    sigmoid_slope_negative = isinstance(sigmoid_slope, float) and sigmoid_slope < 0 # * sigmoid_slope can be array for finite faultskA
-    
-    if LEGACY := True and not sigmoid_slope_negative: # * Here we branch to the experimental activation function with hard sigmoid
-        sigm = activate_formation_block_from_args(Z_x, ids, scalar_value_at_sp, sigmoid_slope)
-    else:
-        from .torch_activation import activate_formation_block_from_args_hard_sigmoid
-        sigm = activate_formation_block_from_args_hard_sigmoid(Z_x, ids, scalar_value_at_sp)
+    sigmoid_slope_negative = isinstance(sigmoid_slope, float) and sigmoid_slope < 0  # * sigmoid_slope can be array for finite faultskA
 
+    if LEGACY := False and not sigmoid_slope_negative:  # * Here we branch to the experimental activation function with hard sigmoid
+        sigm = activate_formation_block_from_args(
+            Z_x=Z_x,
+            ids=ids,
+            scalar_value_at_sp=scalar_value_at_sp,
+            sigmoid_slope=sigmoid_slope
+        )
+    else:
+        sigm = soft_segment_unbounded(
+            Z=Z_x,
+            edges=scalar_value_at_sp,
+            ids=ids,
+            sigmoid_slope=sigmoid_slope
+        )
     return sigm
 
 

gempy_engine/modules/activator/torch_activation.py

@@ -1,3 +1,5 @@
+import warnings
+
 import torch
 from ...core.backend_tensor import BackendTensor as bt, BackendTensor
 
@@ -14,6 +16,9 @@
 
 
 def activate_formation_block_from_args_hard_sigmoid(Z_x, ids, scalar_value_at_sp):
+    
+    warnings.warn(DeprecationWarning("This function is deprecated. Use activate_formation_block instead."))
+    
     element_0 = bt.t.array([0], dtype=BackendTensor.dtype_obj)
 
     min_Z_x = BackendTensor.t.min(Z_x, axis=0).reshape(-1)  # ? Is this as good as it gets?

tests/fixtures/complex_geometries.py

@@ -35,7 +35,7 @@ def one_fault_model():
 
     spi: SurfacePoints = SurfacePoints(sp_coords)
     ori: Orientations = Orientations(dip_postions, dip_gradients)
-    ids = np.array([1, 2, 3, 4, 5, 6])
+    ids = np.array([1, 2, 3, 4, 5, 6, 7])
 
     resolution = [2, 2, 2]
     extent = np.array([-500, 500., -500, 500, -450, 550]) / rescaling_factor
@@ -172,7 +172,7 @@ def graben_fault_model():
 
     spi = SurfacePoints(sp_coords)
     ori = Orientations(dip_postions, dip_gradients)
-    ids = np.array([1, 2, 3, 4, 5, 6])
+    ids = np.array([1, 2, 3, 4, 5, 6, 7])
 
     resolution = [2, 2, 2]
     extent = np.array([-500, 500., -500, 500, -450, 550]) / rescaling_factor

tests/test_common/test_api/test_faults/test_faults_graben.py

@@ -27,7 +27,7 @@ def test_graben_fault_model(graben_fault_model):
     options.evaluation_options.dual_conturing_fancy = True
     options.debug=True
 
-    options.evaluation_options.number_octree_levels = 4
+    options.evaluation_options.number_octree_levels = 5
     solutions: Solutions = compute_model(interpolation_input, options, structure)
 
     outputs: list[OctreeLevel] = solutions.octrees_output

tests/test_common/test_api/test_faults/test_one_fault.py

@@ -21,7 +21,7 @@
 from gempy_engine.plugins.plotting.helper_functions import plot_block_and_input_2d, plot_scalar_and_input_2d
 
 
-def test_one_fault_model(one_fault_model, n_oct_levels=3):
+def test_one_fault_model(one_fault_model, n_oct_levels=5):
     interpolation_input: InterpolationInput
     structure: InputDataDescriptor
     options: InterpolationOptions
@@ -44,7 +44,7 @@ def test_one_fault_model(one_fault_model, n_oct_levels=3):
         gempy_v2_cov = _covariance_for_one_fault_model_from_gempy_v2()
         diff = last_cov - gempy_v2_cov
 
-    if plot_2d := False:
+    if plot_2d := True:
         _plot_stack_raw(interpolation_input, outputs, structure)
         _plot_stack_squeezed_mask(interpolation_input, outputs, structure)
         _plot_stack_mask_component(interpolation_input, outputs, structure)

tests/test_common/test_modules/test_activator_fns.py

@@ -0,0 +1,100 @@
+import dataclasses
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from gempy_engine.API.interp_single._interp_scalar_field import _solve_interpolation, _evaluate_sys_eq
+from gempy_engine.API.interp_single._interp_single_feature import input_preprocess
+from gempy_engine.config import AvailableBackends
+from gempy_engine.core.data.internal_structs import SolverInput
+from gempy_engine.modules.activator.activator_interface import activate_formation_block
+from gempy_engine.core.backend_tensor import BackendTensor
+
+dir_name = os.path.dirname(__file__)
+
+plot = True
+
+
+def test_activator_3_layers_segmentation_function(simple_model_3_layers, simple_grid_3d_more_points_grid):
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.numpy,
+        ids=np.array([1, 20, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
+
+    if plot:
+        _plot_continious(grid, ids_block, interpolation_input)
+
+
+def test_activator_3_layers_segmentation_function_II(simple_model_3_layers, simple_grid_3d_more_points_grid):
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.numpy,
+        ids=np.array([1, 2, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
+
+    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
+
+    if plot:
+        _plot_continious(grid, ids_block, interpolation_input)
+
+
+def test_activator_3_layers_segmentation_function_torch(simple_model_3_layers, simple_grid_3d_more_points_grid):
+    Z_x, grid, ids_block, interpolation_input = _run_test(
+        backend=AvailableBackends.PYTORCH,
+        ids=np.array([1, 2, 3, 4]),
+        simple_grid_3d_more_points_grid=simple_grid_3d_more_points_grid,
+        simple_model_3_layers=simple_model_3_layers
+    )
+
+    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
+    if plot:
+        _plot_continious(grid, ids_block, interpolation_input)
+
+
+def _run_test(backend, ids, simple_grid_3d_more_points_grid, simple_model_3_layers):
+    interpolation_input = simple_model_3_layers[0]
+    options = simple_model_3_layers[1]
+    data_shape = simple_model_3_layers[2].tensors_structure
+    grid = dataclasses.replace(simple_grid_3d_more_points_grid)
+    interpolation_input.set_temp_grid(grid)
+    interp_input: SolverInput = input_preprocess(data_shape, interpolation_input)
+    weights = _solve_interpolation(interp_input, options.kernel_options)
+    exported_fields = _evaluate_sys_eq(interp_input, weights, options)
+    exported_fields.set_structure_values(
+        reference_sp_position=data_shape.reference_sp_position,
+        slice_feature=interpolation_input.slice_feature,
+        grid_size=interpolation_input.grid.len_all_grids)
+    Z_x: np.ndarray = exported_fields.scalar_field
+    sasp = exported_fields.scalar_field_at_surface_points
+    print(Z_x, Z_x.shape[0])
+    print(sasp)
+    BackendTensor.change_backend_gempy(backend)
+    ids_block = activate_formation_block(
+        exported_fields=exported_fields,
+        ids=ids,
+        sigmoid_slope=500 * 4
+    )[0, :-7]
+    return Z_x, grid, ids_block, interpolation_input
+
+
+def _plot_continious(grid, ids_block, interpolation_input):
+    block__ = ids_block[grid.dense_grid_slice]
+    unique = np.unique(block__)
+    t = block__.reshape(50, 5, 50)[:, 2, :].T
+    unique = np.unique(t)
+
+    levels = np.linspace(t.min(), t.max(), 40)
+    plt.contourf(
+        t,
+        levels=levels,
+        cmap="jet",
+        extent=(.25, .75, .25, .75)
+    )
+    xyz = interpolation_input.surface_points.sp_coords
+    plt.plot(xyz[:, 0], xyz[:, 2], "o")
+    plt.colorbar()
+    plt.show()