Implemented an alternative soft clamp.

ppafm/ocl/cl/FF.cl

@@ -1175,8 +1175,8 @@ __kernel void sumSingleGroup(__global float2 *array, int n) {
 
 }
 
-// Clamp array values to specified range
-__kernel void clamp(
+// Clamp array values to specified range using a hard clamp
+__kernel void clamp_hard(
     __global float *array_in,   // Input array
     __global float *array_out,  // Output array
     int n,                      // Number of elements in array
@@ -1188,6 +1188,25 @@ __kernel void clamp(
     array_out[ind] = max(min_value, min(max_value, array_in[ind]));
 }
 
+// Clamp array values to specified range using a soft clamp
+__kernel void clamp_soft(
+    __global float *array_in,   // Input array
+    __global float *array_out,  // Output array
+    int n,                      // Number of elements in array
+    float min_value,            // Minimum clamp value
+    float max_value,            // Maximum clamp value
+    float width                 // Width of transition region
+) {
+    int ind = get_global_id(0);
+    if (ind >= n) return;
+
+    float val = array_in[ind];
+    val = (val - max_value) / (1 + exp((val - max_value) /  width)) + max_value;
+    val = (val - min_value) / (1 + exp((val - min_value) / -width)) + min_value;
+
+    array_out[ind] = val;
+}
+
 // Multiply and add values in one array with another
 __kernel void addMult(
     __global float *array_in1,  // Input array 1

ppafm/ocl/field.py

@@ -287,13 +287,16 @@ def _prepare_same_size_output_grid(self, array_in, in_place):
             grid_out = array_type(array_out, lvec=self.lvec, shape=self.shape, ctx=self.ctx)
         return grid_out
 
-    def clamp(self, minimum=-np.inf, maximum=np.inf, in_place=True, local_size=(32,), queue=None):
+    def clamp(self, minimum=-np.inf, maximum=np.inf, clamp_type="hard", soft_clamp_width=1.0, in_place=True, local_size=(32,), queue=None):
         """
-        Clamp data grid values to a specified range.
+        Clamp data grid values to a specified range. The ``'hard'`` clamp simply clips values that are out of range,
+        and the ``'soft'`` clamp uses a sigmoid to smoothen the transition.
 
         Arguments:
             minimum: float. Values below minimum are set to minimum.
-            maximum: float. Values above maximum are set to maximum
+            maximum: float. Values above maximum are set to maximum.
+            clamp_type: str. Type of clamp to use: ``'soft'`` or ``'hard'``.
+            soft_clamp_width: float. Width of transition region for soft clamp.
             in_place: bool. Whether to do operation in place or to create a new array.
             local_size: tuple of a single int. Size of local work group on device.
             queue: pyopencl.CommandQueue. OpenCL queue on which operation is performed. Defaults to oclu.queue.
@@ -307,19 +310,34 @@ def clamp(self, minimum=-np.inf, maximum=np.inf, in_place=True, local_size=(32,)
         n = np.int32(array_in.size / 4)
         minimum = np.float32(minimum)
         maximum = np.float32(maximum)
+        soft_clamp_width = np.float32(soft_clamp_width)
 
         queue = queue or oclu.queue
         global_size = [int(np.ceil(n / local_size[0]) * local_size[0])]
 
-        # fmt: off
-        cl_program.clamp(queue, global_size, local_size,
-            array_in,
-            grid_out.cl_array,
-            n,
-            minimum,
-            maximum,
-        )
-        # fmt: on
+        if clamp_type == "hard":
+            # fmt: off
+            cl_program.clamp_hard(queue, global_size, local_size,
+                array_in,
+                grid_out.cl_array,
+                n,
+                minimum,
+                maximum,
+            )
+            # fmt: on
+        elif clamp_type == "soft":
+            # fmt: off
+            cl_program.clamp_soft(queue, global_size, local_size,
+                array_in,
+                grid_out.cl_array,
+                n,
+                minimum,
+                maximum,
+                soft_clamp_width,
+            )
+            # fmt: on
+        else:
+            raise ValueError(f"Unsupported clamp type `{clamp_type}`")
 
         return grid_out
 

tests/human_eye/test_clamp.py

@@ -0,0 +1,31 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+import ppafm.ocl.oclUtils as oclu
+from ppafm.ocl.field import DataGrid
+
+oclu.init_env(i_platform=0)
+
+
+def test_clamp_visual():
+
+    triangle_np = np.concatenate([np.linspace(-3, 10, 100), np.linspace(10, -3, 100)[1:]])
+    triangle_data_grid = DataGrid(triangle_np[None, None], lvec=np.concatenate([np.zeros((1, 3)), np.eye(3)], axis=0))
+
+    minimum = -1
+    maximum = 5
+    width = 0.5
+    triangle_clamp_hard = triangle_data_grid.clamp(minimum=minimum, maximum=maximum, clamp_type="hard", in_place=False).array[0, 0]
+    triangle_clamp_soft = triangle_data_grid.clamp(minimum=minimum, maximum=maximum, clamp_type="soft", soft_clamp_width=width, in_place=False).array[0, 0]
+
+    plt.plot(triangle_np)
+    plt.plot(triangle_clamp_hard)
+    plt.plot(triangle_clamp_soft)
+
+    plt.legend(["Original", "Hard clamp", "Soft clamp"])
+    plt.savefig("clamps.png")
+    plt.show()
+
+
+if __name__ == "__main__":
+    test_clamp_visual()

tests/test_datagrid.py

@@ -3,7 +3,6 @@
 import time
 
 import numpy as np
-import pyopencl as cl
 
 import ppafm.ocl.field as FFcl
 import ppafm.ocl.oclUtils as oclu