[ENH] Adding numpy version of activation function

Author: Leguark

gempy_engine/modules/activator/activator_interface.py

@@ -15,14 +15,28 @@ def activate_formation_block(exported_fields: ExportedFields, ids: np.ndarray,
     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, ids, scalar_value_at_sp, sigmoid_slope)
-    else:
-        sigm = soft_segment_unbounded(
-            Z=Z_x,
-            edges=scalar_value_at_sp,
-            ids=ids,
+        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:
+        match BackendTensor.engine_backend:
+            case AvailableBackends.PYTORCH:
+                sigm = soft_segment_unbounded(
+                    Z=Z_x,
+                    edges=scalar_value_at_sp,
+                    ids=ids,
+                    sigmoid_slope=sigmoid_slope
+                )
+            case AvailableBackends.numpy:
+                sigm = soft_segment_unbounded_np(
+                    Z=Z_x,
+                    edges=scalar_value_at_sp,
+                    ids=ids,
+                    sigmoid_slope=sigmoid_slope   
+                )
 
     return sigm
 
@@ -110,3 +124,43 @@ def soft_segment_unbounded(Z, edges, ids, sigmoid_slope):
     # weighted sum by the ids
     ids__sum = (membership * ids).sum(dim=-1)
     return np.atleast_2d(ids__sum.numpy())
+
+
+import numpy as np
+
+def soft_segment_unbounded_np(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
+    """
+    Z     = np.asarray(Z)
+    edges = np.asarray(edges)
+    ids   = np.asarray(ids)
+
+    # 1) per-edge temperatures τ_k = |Δ_k|/(4·m)
+    jumps = np.abs(ids[1:] - ids[:-1])      # shape (K-1,)
+    tau_k = jumps / (4 * sigmoid_slope)     # shape (K-1,)
+
+    # 2) first bin (-∞, e1) via σ((e1 - Z)/τ₁)
+    first = 1.0 / (1.0 + np.exp((Z - edges[0]) / 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 (...,)
+
+    # 4) middle bins [e_i, e_{i+1}): σ((Z - e_i)/τ_i) - σ((Z - e_{i+1})/τ_{i+1})
+    Z_exp = Z[..., None]  # shape (...,1)
+    left  =  1.0 / (1.0 + np.exp(-(Z_exp - edges[:-1]) / tau_k[:-1]))  # (...,K-2)
+    right =  1.0 / (1.0 + np.exp(-(Z_exp - edges[1: ]) / tau_k[1: ]))  # (...,K-2)
+    middle = left - right                                              # (...,K-2)
+
+    # 5) assemble memberships and weight by ids
+    membership = np.concatenate(
+        [ first[..., None], middle, last[..., None] ],
+        axis=-1
+    )  # shape (...,K)
+
+    ids__sum = np.sum(membership * ids, axis=-1)
+    return np.atleast_2d(ids__sum)

tests/test_common/test_modules/test_activator_fns.py

@@ -41,25 +41,18 @@ def test_activator_3_layers_segmentation_function(simple_model_3_layers, simple_
     print(Z_x, Z_x.shape[0])
     print(sasp)
 
-    BackendTensor.change_backend_gempy(AvailableBackends.PYTORCH)
     ids_block = activate_formation_block(
         exported_fields=exported_fields,
         ids=ids,
-        # sigmoid_slope=5_000_000
         sigmoid_slope = 500*4
     )[0, :-7]
 
-    BackendTensor.change_backend_gempy(AvailableBackends.numpy)
-
     if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
         ids_block = ids_block.detach().numpy()
         Z_x = Z_x.detach().numpy()
         interpolation_input.surface_points.sp_coords = interpolation_input.surface_points.sp_coords.detach().numpy()
 
     if plot:
-        if False:
-            _plot_4_categories(grid, ids_block, interpolation_input)
-
         _plot_continious(grid, ids_block, interpolation_input)
 
 
@@ -82,22 +75,3 @@ def _plot_continious(grid, ids_block, interpolation_input):
     plt.show()
 
 
-def _plot_4_categories(grid, ids_block, interpolation_input):
-    block__ = ids_block[grid.dense_grid_slice]
-    reshaped_block = block__.reshape(50, 5, 50)[:, 2, :].T
-    bounds = [0, 1, 2, 3, 4]
-    # Create a normalized colormap with 5 distinct colors
-    import matplotlib.colors as colors
-    cmap = plt.cm.autumn
-    norm = colors.BoundaryNorm(bounds, cmap.N)
-    # Plot with 5 distinct categories
-    im = plt.contourf(reshaped_block, levels=bounds, cmap=cmap, norm=norm,
-                      extent=(.25, .75, .25, .75))
-    xyz = interpolation_input.surface_points.sp_coords
-    plt.plot(xyz[:, 0], xyz[:, 2], "o")
-    # Add a colorbar with category labels
-    cbar = plt.colorbar(im, ticks=bounds)
-    # cbar.ax.set_yticklabels(['', '<1', '1-2', '2-3', '3-4', '>4'])
-    cbar.ax.set_yticklabels(['', '<1', '1-2', '2-3', '3-4'])
-    plt.title('Formation Categories')
-    plt.show()