[ENH] Implementing soft_segment_unbounded for numpy and torch

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,12 +1,13 @@
 import warnings
 
-from gempy_engine.config import DEBUG_MODE, AvailableBackends
-from gempy_engine.core.backend_tensor import BackendTensor as bt, BackendTensor
+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
+
 import numpy as np
 import numbers
 
-from gempy_engine.core.data.exported_fields import ExportedFields
-
 
 def activate_formation_block(exported_fields: ExportedFields, ids: np.ndarray,
                              sigmoid_slope: float) -> np.ndarray:
@@ -23,9 +24,17 @@ def activate_formation_block(exported_fields: ExportedFields, ids: np.ndarray,
             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
+
         match BackendTensor.engine_backend:
             case AvailableBackends.PYTORCH:
-                sigm = soft_segment_unbounded(
+                sigm = soft_segment_unbounded_torch(
                     Z=Z_x,
                     edges=scalar_value_at_sp,
                     ids=ids,
@@ -85,7 +94,7 @@ def _compute_sigmoid(Z_x, scale_0, scale_1, drift_0, drift_1, drift_id, sigmoid_
 import torch
 
 
-def soft_segment_unbounded(Z, edges, ids, sigmoid_slope):
+def soft_segment_unbounded_torch(Z, edges, ids, sigmoid_slope):
     """
     Z:            (...,) tensor of scalar values
     edges:   (K-1,) tensor of finite split points [e1, e2, ..., e_{K-1}]
@@ -124,7 +133,7 @@ def soft_segment_unbounded(Z, edges, ids, sigmoid_slope):
 
     # weighted sum by the ids
     ids__sum = (membership * ids).sum(dim=-1)
-    
+
     # make it at least 2d
     ids__sum = ids__sum[None, :]
 
@@ -151,10 +160,9 @@ def soft_segment_unbounded_np(Z, edges, ids, sigmoid_slope):
         case np.ndarray():
             membership = _final_faults_segmentation(Z, edges, sigmoid_slope)
         case numbers.Number():
-             membership = _lith_segmentation(Z, edges, ids, sigmoid_slope)
+            membership = _lith_segmentation(Z, edges, ids, sigmoid_slope)
         case _:
-             raise ValueError("sigmoid_slope must be a float or an array")
-
+            raise ValueError("sigmoid_slope must be a float or an array")
 
     ids__sum = np.sum(membership * ids, axis=-1)
     return np.atleast_2d(ids__sum)
@@ -179,7 +187,6 @@ def _final_faults_segmentation(Z, edges, sigmoid_slope):
 
 
 def _lith_segmentation(Z, edges, ids, sigmoid_slope):
-    
     # 1) per-edge temperatures τ_k = |Δ_k|/(4·m)
     jumps = np.abs(ids[1:] - ids[:-1])  # shape (K-1,)
     tau_k = jumps / float(sigmoid_slope)  # shape (K-1,)