From d7123ccc1cbff79f0a124895f10c1a2f16df7416 Mon Sep 17 00:00:00 2001 From: NikoOinonen Date: Mon, 14 Oct 2024 10:48:44 +0300 Subject: [PATCH] Add clamp function to DataGrid. --- ppafm/ocl/cl/FF.cl | 13 +++++++++++++ ppafm/ocl/field.py | 39 +++++++++++++++++++++++++++++++++++++++ tests/test_datagrid.py | 15 +++++++++++++++ 3 files changed, 67 insertions(+) diff --git a/ppafm/ocl/cl/FF.cl b/ppafm/ocl/cl/FF.cl index 51ab32ae..b5be1fba 100644 --- a/ppafm/ocl/cl/FF.cl +++ b/ppafm/ocl/cl/FF.cl @@ -1175,6 +1175,19 @@ __kernel void sumSingleGroup(__global float2 *array, int n) { } +// Clamp array values to specified range +__kernel void clamp( + __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 +) { + int ind = get_global_id(0); + if (ind >= n) return; + array_out[ind] = max(min_value, min(max_value, array_in[ind])); +} + // Multiply and add values in one array with another __kernel void addMult( __global float *array_in1, // Input array 1 diff --git a/ppafm/ocl/field.py b/ppafm/ocl/field.py index b0b37a08..c8ed0c1f 100644 --- a/ppafm/ocl/field.py +++ b/ppafm/ocl/field.py @@ -277,6 +277,45 @@ def to_file(self, file_path, clamp=None): io.save_vec_field(file_head, array[:, :, :, :3], self.lvec, data_format=ext) io.save_scal_field(file_head + "_w", array[:, :, :, 3], self.lvec, data_format=ext) + def clamp(self, minimum=-np.inf, maximum=np.inf, in_place=True, local_size=(32,), queue=None): + """ + Clamp data grid values to a specified range. + + Arguments: + minimum: float. Values below minimum are set to minimum. + maximum: float. Values above maximum are set to maximum + 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. + + Returns: + grid_out: Same type as self. New data grid with result. + """ + array_in = self.cl_array + queue = queue or oclu.queue + if in_place: + grid_out = self + array_out = array_in + self._array = None # The current numpy array will be wrong after operation so reset it + else: + array_out = cl.Buffer(self.ctx, cl.mem_flags.READ_WRITE, size=array_in.size) + array_type = type(self) # This way so inherited classes return their own class type + grid_out = array_type(array_out, lvec=self.lvec, shape=self.shape, ctx=self.ctx) + n = np.int32(array_in.size / 4) + minimum = np.float32(minimum) + maximum = np.float32(maximum) + global_size = [int(np.ceil(n / local_size[0]) * local_size[0])] + # fmt: off + cl_program.clamp(queue, global_size, local_size, + array_in, + array_out, + n, + minimum, + maximum, + ) + # fmt: on + return grid_out + def add_mult(self, array, scale=1.0, in_place=True, local_size=(32,), queue=None): """ Multiply the values of another data grid and add them to the values of this data grid. diff --git a/tests/test_datagrid.py b/tests/test_datagrid.py index 49de4085..8096be42 100755 --- a/tests/test_datagrid.py +++ b/tests/test_datagrid.py @@ -70,3 +70,18 @@ def test_tip_interp(): # Check that the interpolant is close to the analytical solution assert np.allclose(rho_interp.array, rho2.array, atol=1e-3, rtol=1e-3) + + +def test_clamp(): + + array = np.array([-1.0, 0.0, 1.0, 2.0])[None, None] + data_grid = FFcl.DataGrid(array, lvec=np.concatenate([np.zeros((1, 3)), np.eye(3)], axis=0)) + + new_grid = data_grid.clamp(minimum=0.0, maximum=1.0, in_place=False) + + assert np.allclose(new_grid.array, [0.0, 0.0, 1.0, 1.0]) + assert np.allclose(data_grid.array, array) + + data_grid.clamp(maximum=1.0, in_place=True) + + assert np.allclose(data_grid.array, [-1.0, 0.0, 1.0, 1.0])