nabu 2023.2.1__py3-none-any.whl → 2024.1.0rc3__py3-none-any.whl

nabu/processing/tests/test_fft.py

@@ -0,0 +1,268 @@
+from itertools import permutations
+import pytest
+import numpy as np
+from scipy.fft import fftn, ifftn, rfftn, irfftn
+from nabu.testutils import generate_tests_scenarios, get_data, get_array_of_given_shape, __do_long_tests__
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.processing.fft_cuda import SKCUFFT, VKCUFFT, has_vkfft as has_cuda_vkfft
+from nabu.opencl.utils import __has_pyopencl__, get_opencl_context
+from nabu.processing.fft_opencl import VKCLFFT, has_vkfft as has_cl_vkfft
+from nabu.processing.fft_base import is_fast_axes
+
+try:
+    import skcuda
+
+    __has_skcuda__ = True
+except ImportError:
+    __has_skcuda__ = False
+
+
+scenarios = {
+    "shape": [(256,), (300,), (300, 301), (300, 302)],
+    "r2c": [True, False],
+    "precision": ["simple"],
+    "backend": ["cuda", "opencl"],
+}
+
+if __do_long_tests__:
+    scenarios["shape"].extend([(307,), (125, 126, 260)])
+    scenarios["precision"].append("double")
+
+scenarios = generate_tests_scenarios(scenarios)
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+
+    cls.data = get_data("chelsea.npz")["data"]
+    cls.abs_tol = {
+        "simple": {
+            1: 5e-3,
+            2: 1.0e0,
+            3: 5e2,  # !
+        },
+        "double": {
+            1: 1e-10,
+            2: 1e-9,
+            3: 1e-7,
+        },
+    }
+    if __has_pycuda__:
+        cls.cu_ctx = get_cuda_context(cleanup_at_exit=False)
+    if __has_pyopencl__:
+        cls.cl_ctx = get_opencl_context("all")
+    yield
+    if __has_pycuda__:
+        cls.cu_ctx.pop()
+
+
+def _get_fft_cls(backend):
+    fft_cls = None
+    if backend == "cuda":
+        if not (has_cuda_vkfft() and __has_pycuda__):
+            pytest.skip("Need vkfft and pycuda to use VKCUFFT")
+        fft_cls = VKCUFFT
+    if backend == "opencl":
+        if not (has_cl_vkfft() and __has_pyopencl__):
+            pytest.skip("Need vkfft and pyopencl to use VKCLFFT")
+        fft_cls = VKCLFFT
+    return fft_cls
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestFFT:
+    def _get_data_array(self, config):
+        r2c = config["r2c"]
+        shape = config["shape"]
+        precision = config["precision"]
+        dtype = {
+            True: {"simple": np.float32, "double": np.float64},
+            False: {"simple": np.complex64, "double": np.complex128},
+        }[r2c][precision]
+        data = get_array_of_given_shape(self.data, shape, dtype)
+        return data
+
+    @staticmethod
+    def check_result(res, ref, config, tol, name=""):
+        err_max = np.max(np.abs(res - ref))
+        err_msg = "%s FFT(%s, r2c=%s): tol=%.2e, but max error = %.2e" % (
+            name,
+            str(config["shape"]),
+            str(config["r2c"]),
+            tol,
+            err_max,
+        )
+        assert np.allclose(res, ref, atol=tol), err_msg
+
+    def _do_fft(self, data, r2c, axes=None, return_fft_obj=False, backend_cls=None):
+        ctx = self.cu_ctx if backend_cls.backend == "cuda" else self.cl_ctx
+        fft = backend_cls(data.shape, data.dtype, r2c=r2c, axes=axes, ctx=ctx)
+        d_data = fft.processing.allocate_array("_data", data.shape, dtype=data.dtype)
+        d_data.set(data)
+        d_out = fft.fft(d_data)
+        res = d_out.get()
+        return (res, fft) if return_fft_obj else res
+
+    @staticmethod
+    def _do_reference_fft(data, r2c, axes=None):
+        ref_fft_func = rfftn if r2c else fftn
+        ref = ref_fft_func(data, axes=axes)
+        return ref
+
+    @staticmethod
+    def _do_reference_ifft(data, r2c, axes=None):
+        ref_ifft_func = irfftn if r2c else ifftn
+        ref = ref_ifft_func(data, axes=axes)
+        return ref
+
+    @pytest.mark.skipif(not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda scikit-cuda for this test")
+    @pytest.mark.parametrize("config", scenarios)
+    def test_sckcuda(self, config):
+        r2c = config["r2c"]
+        shape = config["shape"]
+        precision = config["precision"]
+        ndim = len(shape)
+        if ndim == 3 and not (__do_long_tests__):
+            pytest.skip("3D FFTs are done only for long tests - use NABU_LONG_TESTS=1")
+
+        data = self._get_data_array(config)
+
+        res, cufft = self._do_fft(data, r2c, return_fft_obj=True, backend_cls=SKCUFFT)
+        ref = self._do_reference_fft(data, r2c)
+
+        tol = self.abs_tol[precision][ndim]
+        self.check_result(res, ref, config, tol, name="skcuda")
+
+        # Complex-to-complex can also be performed on real data (as in numpy.fft.fft(real_data))
+        if not (r2c):
+            res = self._do_fft(data, False, backend_cls=SKCUFFT)
+            ref = self._do_reference_fft(data, False)
+            self.check_result(res, ref, config, tol, name="skcuda")
+
+        # IFFT
+        res = cufft.ifft(cufft.output_fft).get()
+        self.check_result(res, data, config, tol, name="skcuda")
+        # Perhaps we should also check against numpy/scipy ifft,
+        # but it does not yield the good shape for R2C on odd-sized data
+
+    @pytest.mark.skipif(not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda scikit-cuda for this test")
+    @pytest.mark.parametrize("config", scenarios)
+    def test_skcuda_batched(self, config):
+        shape = config["shape"]
+        if len(shape) == 1:
+            return
+        elif len(shape) == 3 and not (__do_long_tests__):
+            pytest.skip("3D FFTs are done only for long tests - use NABU_LONG_TESTS=1")
+        r2c = config["r2c"]
+        tol = self.abs_tol[config["precision"]][len(shape)]
+
+        data = self._get_data_array(config)
+
+        if data.ndim == 2:
+            axes_to_test = [(0,), (1,)]
+        elif data.ndim == 3:
+            # axes_to_test = [(1, 2), (2, 1), (2,)] # See fft.py: works for C2C but not R2C ?
+            axes_to_test = [(2,)]
+
+        for axes in axes_to_test:
+            res, cufft = self._do_fft(data, r2c, axes=axes, return_fft_obj=True, backend_cls=SKCUFFT)
+            ref = self._do_reference_fft(data, r2c, axes=axes)
+            self.check_result(res, ref, config, tol, name="skcuda batched axes=%s" % (str(axes)))
+            # IFFT
+            res = cufft.ifft(cufft.output_fft).get()
+            self.check_result(res, data, config, tol, name="skcuda")
+
+    @pytest.mark.parametrize("config", scenarios)
+    def test_vkfft(self, config):
+        backend = config["backend"]
+        fft_cls = _get_fft_cls(backend)
+
+        r2c = config["r2c"]
+        shape = config["shape"]
+        precision = config["precision"]
+        ndim = len(shape)
+        if ndim == 3 and not (__do_long_tests__):
+            pytest.skip("3D FFTs are done only for long tests - use NABU_LONG_TESTS=1")
+
+        if ndim >= 2 and r2c and shape[-1] & 1:
+            pytest.skip("R2C with odd-sized fast dimension is not supported in VKFFT")
+
+        data = self._get_data_array(config)
+
+        res, fft_obj = self._do_fft(data, r2c, return_fft_obj=True, backend_cls=fft_cls)
+        ref = self._do_reference_fft(data, r2c)
+
+        tol = self.abs_tol[precision][ndim]
+        self.check_result(res, ref, config, tol, name="vkfft_%s" % backend)
+
+        # Complex-to-complex can also be performed on real data (as in numpy.fft.fft(real_data))
+        if not (r2c):
+            res = self._do_fft(data, False, backend_cls=fft_cls)
+            ref = self._do_reference_fft(data, False)
+            self.check_result(res, ref, config, tol, name="vkfft_%s" % backend)
+
+        # IFFT
+        res = fft_obj.ifft(fft_obj.output_fft).get()
+        self.check_result(res, data, config, tol, name="vkfft_%s" % backend)
+
+    @pytest.mark.parametrize("config", scenarios)
+    def test_vkfft_batched(self, config):
+        backend = config["backend"]
+        fft_cls = _get_fft_cls(backend)
+        shape = config["shape"]
+        if len(shape) == 1:
+            return
+        elif len(shape) == 3 and not (__do_long_tests__):
+            pytest.skip("3D FFTs are done only for long tests - use NABU_LONG_TESTS=1")
+        r2c = config["r2c"]
+        tol = self.abs_tol[config["precision"]][len(shape)]
+
+        data = self._get_data_array(config)
+
+        if data.ndim >= 2 and r2c and shape[-1] & 1:
+            pytest.skip("R2C with odd-sized fast dimension is not supported in VKFFT")
+
+        # For R2C, only fastest axes are supported by vkfft
+        if data.ndim == 2:
+            axes_to_test = [(1,)]
+        elif data.ndim == 3:
+            axes_to_test = [
+                (1, 2),
+                (2,),
+            ]
+        for axes in axes_to_test:
+            res, cufft = self._do_fft(data, r2c, axes=axes, return_fft_obj=True, backend_cls=fft_cls)
+            ref = self._do_reference_fft(data, r2c, axes=axes)
+            self.check_result(res, ref, config, tol, name="vkfft_%s batched axes=%s" % (backend, str(axes)))
+            # IFFT
+            res = cufft.ifft(cufft.output_fft).get()
+            self.check_result(res, data, config, tol, name="vkfft_%s" % backend)
+
+    @pytest.mark.skipif(not (__do_long_tests__), reason="Use NABU_LONG_TESTS=1 for this test")
+    def test_fast_axes_utility_function(self):
+        axes_to_test = {
+            2: {
+                (0, 1): True,
+                (1,): True,
+                (-1,): True,
+                (-2,): False,
+                (0,): False,
+            },
+            3: {
+                (0, 1, 2): True,
+                (0, 1): False,
+                (1, 2): True,
+                (2, 1): True,
+                (-2, -1): True,
+                (2,): True,
+                (-1,): True,
+            },
+        }
+        for ndim, axes_ in axes_to_test.items():
+            for axes, is_fast in axes_.items():
+                possible_axes = [axes]
+                if len(axes) > 1:
+                    possible_axes = list(permutations(axes, len(axes)))
+                for ax in possible_axes:
+                    assert is_fast_axes(ndim, ax) is is_fast

nabu/processing/tests/test_fftshift.py

@@ -0,0 +1,71 @@
+import numpy as np
+import pytest
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.opencl.utils import __has_pyopencl__, get_opencl_context
+from nabu.testutils import get_data, generate_tests_scenarios, __do_long_tests__
+
+if __has_pyopencl__:
+    from nabu.processing.fftshift import OpenCLFFTshift
+
+configs = {
+    "shape": [(300, 451), (300, 300), (255, 300)],
+    "axes": [(1,)],
+    "dtype_in_out": [(np.float32, np.complex64), (np.complex64, np.float32)],
+    "inplace": [True, False],
+}
+
+scenarios = generate_tests_scenarios(configs)
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.data = get_data("chelsea.npz")["data"]
+    cls.tol = 1e-7
+    if __has_pycuda__:
+        cls.cu_ctx = get_cuda_context(cleanup_at_exit=False)
+    if __has_pyopencl__:
+        cls.cl_ctx = get_opencl_context(device_type="all")
+    yield
+    if __has_pycuda__:
+        cls.cu_ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestFFTshift:
+    def _do_test_fftshift(self, config, fftshift_cls):
+        shape = config["shape"]
+        dtype = config["dtype_in_out"][0]
+        dst_dtype = config["dtype_in_out"][1]
+        axes = config["axes"]
+        inplace = config["inplace"]
+        if inplace and shape[-1] & 1:
+            pytest.skip("Not Implemented")
+        data = np.ascontiguousarray(self.data[: shape[0], : shape[1]], dtype=dtype)
+
+        backend = fftshift_cls.backend
+        ctx = self.cu_ctx if backend == "cuda" else self.cl_ctx
+        backend_options = {"ctx": ctx}
+        if not (inplace):
+            fftshift = fftshift_cls(data.shape, dtype, dst_dtype=dst_dtype, axes=axes, **backend_options)
+        else:
+            fftshift = fftshift_cls(data.shape, dtype, axes=axes, **backend_options)
+
+        d_data = fftshift.processing.allocate_array("data", shape, dtype)
+        d_data.set(data)
+
+        d_res = fftshift.fftshift(d_data)
+
+        assert (
+            np.max(np.abs(d_res.get() - np.fft.fftshift(data, axes=axes))) == 0
+        ), "something wrong with fftshift_%s(%s)" % (backend, str(config))
+
+    # @pytest.mark.skipif(not (__has_pycuda__), reason="Need pycuda for this test")
+    # @pytest.mark.parametrize("config", scenarios)
+    # def test_cuda_transpose(self, config):
+    # self._do_test_transpose(config, CudaTranspose)
+
+    @pytest.mark.skipif(not (__has_pyopencl__), reason="Need pyopencl for this test")
+    @pytest.mark.parametrize("config", scenarios)
+    def test_opencl_fftshift(self, config):
+        self._do_test_fftshift(config, OpenCLFFTshift)

nabu/{misc → processing}/tests/test_histogram.py

@@ -1,13 +1,11 @@
-from os import path
-from tempfile import mkdtemp
 import pytest
 import numpy as np
 from nabu.testutils import get_data
-from nabu.misc.histogram import PartialHistogram
+from nabu.processing.histogram import PartialHistogram
 from nabu.cuda.utils import __has_pycuda__, get_cuda_context
 
 if __has_pycuda__:
-    from nabu.misc.histogram_cuda import CudaPartialHistogram
+    from nabu.processing.histogram_cuda import CudaPartialHistogram
     import pycuda.gpuarray as garray
 
 

nabu/{cuda → processing}/tests/test_medfilt.py

@@ -5,7 +5,7 @@ from nabu.testutils import generate_tests_scenarios, get_data
 from nabu.cuda.utils import get_cuda_context, __has_pycuda__
 
 if __has_pycuda__:
-    from nabu.cuda.medfilt import MedianFilter
+    from nabu.processing.medfilt_cuda import MedianFilter
     import pycuda.gpuarray as garray
 
 

nabu/processing/tests/test_muladd.py

@@ -0,0 +1,54 @@
+import pytest
+import numpy as np
+from nabu.processing.muladd import MulAdd
+from nabu.testutils import get_data
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+
+if __has_pycuda__:
+    from nabu.processing.muladd_cuda import CudaMulAdd
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.data = get_data("chelsea.npz")["data"].astype("f")  # (300, 451)
+    cls.tol = 1e-7
+    if __has_pycuda__:
+        cls.cu_ctx = get_cuda_context(cleanup_at_exit=False)
+    yield
+    if __has_pycuda__:
+        cls.cu_ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestMulad:
+    def test_muladd(self):
+        dst = self.data.copy()
+        other = self.data.copy()
+        mul_add = MulAdd()
+
+        # Test with no subregion
+        mul_add(dst, other, 1, 2)
+        assert np.allclose(dst, self.data * 1 + other * 2)
+
+        # Test with x-y subregion
+        dst = self.data.copy()
+        mul_add(dst, other, 0.5, 1.7, (slice(10, 200), slice(15, 124)), (slice(100, 290), slice(200, 309)))
+        assert np.allclose(dst[10:200, 15:124], self.data[10:200, 15:124] * 0.5 + self.data[100:290, 200:309] * 1.7)
+
+    @pytest.mark.skipif(not (__has_pycuda__), reason="Need Cuda/pycuda for this test")
+    def test_cuda_muladd(self):
+        mul_add = CudaMulAdd(ctx=self.cu_ctx)
+        dst = mul_add.processing.to_device("dst", self.data)
+        other = mul_add.processing.to_device("other", (self.data / 2).astype("f"))
+
+        # Test with no subregion
+        mul_add(dst, other, 3, 5)
+        assert np.allclose(dst.get(), self.data * 3 + (self.data / 2) * 5)
+
+        # Test with x-y subregion
+        dst.set(self.data)
+        mul_add(dst, other, 0.5, 1.7, (slice(10, 200), slice(15, 124)), (slice(100, 290), slice(200, 309)))
+        assert np.allclose(
+            dst.get()[10:200, 15:124], self.data[10:200, 15:124] * 0.5 + (self.data / 2)[100:290, 200:309] * 1.7
+        )

nabu/{cuda → processing}/tests/test_padding.py

@@ -1,13 +1,130 @@
 import numpy as np
 import pytest
 from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.opencl.utils import __has_pyopencl__, get_opencl_context
+from nabu.processing.padding_cuda import CudaPadding
+from nabu.processing.padding_opencl import OpenCLPadding
 from nabu.utils import calc_padding_lengths, get_cuda_srcfile
+from nabu.testutils import __do_long_tests__
 from nabu.testutils import get_data, generate_tests_scenarios
 
+scenarios = {
+    "shape": [(511, 512), (512, 511)],
+    "pad_width": [((256, 255), (128, 127))],
+    "mode_cuda": CudaPadding.supported_modes[:2] if __has_pycuda__ else [],
+    "mode_opencl": OpenCLPadding.supported_modes[:2] if __has_pyopencl__ else [],
+    "constant_values": [0, ((1.0, 2.0), (3.0, 4.0))],
+    "output_is_none": [True, False],
+    "backend": ["cuda", "opencl"],
+}
+
+
+if __do_long_tests__:
+    scenarios["mode_cuda"] = CudaPadding.supported_modes if __has_pycuda__ else []
+    scenarios["mode_opencl"] = OpenCLPadding.supported_modes if __has_pyopencl__ else []
+    scenarios["pad_width"].extend([((0, 0), (6, 7))])
+
+scenarios = generate_tests_scenarios(scenarios)
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.data = get_data("brain_phantom.npz")["data"]
+    cls.tol = 1e-7
+    if __has_pycuda__:
+        cls.cu_ctx = get_cuda_context(cleanup_at_exit=False)
+    if __has_pyopencl__:
+        cls.cl_ctx = get_opencl_context(device_type="all")
+    yield
+    if __has_pycuda__:
+        cls.cu_ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestPadding:
+    @pytest.mark.parametrize("config", scenarios)
+    def test_padding(self, config):
+        backend = config["backend"]
+        shape = config["shape"]
+        padding_mode = config["mode_%s" % backend]
+        data = self.data[: shape[0], : shape[1]]
+        kwargs = {}
+        if padding_mode == "constant":
+            kwargs["constant_values"] = config["constant_values"]
+        ref = np.pad(data, config["pad_width"], mode=padding_mode, **kwargs)
+
+        PaddingCls = CudaPadding if backend == "cuda" else OpenCLPadding
+        if backend == "cuda":
+            backend_options = {"cuda_options": {"ctx": self.cu_ctx}}
+        else:
+            backend_options = {"opencl_options": {"ctx": self.cl_ctx}}
+
+        padding = PaddingCls(
+            config["shape"],
+            config["pad_width"],
+            mode=padding_mode,
+            constant_values=config["constant_values"],
+            **backend_options,
+        )
+        if config["output_is_none"]:
+            output = None
+        else:
+            output = padding.processing.allocate_array("output", ref.shape, dtype="f")
+
+        d_img = padding.processing.allocate_array("d_img", data.shape, dtype="f")
+        d_img.set(np.ascontiguousarray(data, dtype="f"))
+        res = padding.pad(d_img, output=output)
+
+        err_max = np.max(np.abs(res.get() - ref))
+        assert err_max < self.tol, str("Something wrong with padding for configuration %s" % (str(config)))
+
+    @pytest.mark.skipif(not (__has_pycuda__) and not (__has_pyopencl__), reason="need pycuda or pyopencl")
+    def test_custom_coordinate_transform(self):
+        data = self.data
+        R, C = np.indices(data.shape, dtype=np.int32)
+
+        pad_width = ((256, 255), (254, 251))
+        mode = "reflect"
+
+        coords_R = np.pad(R, pad_width[0], mode=mode)[:, 0]
+        coords_C = np.pad(C, pad_width[1], mode=mode)[0, :]
+
+        # Further transform of coordinates - here FFT layout
+        coords_R = np.roll(coords_R, -pad_width[0][0])
+        coords_C = np.roll(coords_C, -pad_width[1][0])
+
+        padding_classes_to_test = []
+        if __has_pycuda__:
+            padding_classes_to_test.append(CudaPadding)
+        if __has_pyopencl__:
+            padding_classes_to_test.append(OpenCLPadding)
+
+        for padding_cls in padding_classes_to_test:
+            ctx = self.cl_ctx if padding_cls.backend == "opencl" else self.cu_ctx
+            padding = padding_cls(data.shape, (coords_R, coords_C), mode=mode, ctx=ctx)
+
+            d_img = padding.processing.allocate_array("d_img", data.shape, dtype="f")
+            d_img.set(data)
+            d_out = padding.processing.allocate_array("d_out", padding.padded_shape, dtype="f")
+
+            res = padding.pad(d_img, output=d_out)
+
+            ref = np.roll(np.pad(data, pad_width, mode=mode), (-pad_width[0][0], -pad_width[1][0]), axis=(0, 1))
+
+            err_max = np.max(np.abs(d_out.get() - ref))
+            assert err_max < self.tol, "Something wrong with custom padding"
+
+
+#
+# The following is testing a previous version of padding kernels
+# They use specific code (instead of a generic coordinate transform)
+#
+
 if __has_pycuda__:
-    import pycuda.gpuarray as garray
     from nabu.cuda.kernel import CudaKernel
-    from nabu.cuda.padding import CudaPadding
+    import pycuda.gpuarray as garray
+
 
 scenarios_legacy = [
     {
@@ -128,76 +245,3 @@ class TestPaddingLegacy:
         # Compare
         errmax = np.max(np.abs(self.d_data_padded.get() - data_padded_ref))
         assert errmax < self.tol, "Max error is too high"
-
-
-scenarios = generate_tests_scenarios(
-    {
-        "shape": [(511, 512), (512, 511)],
-        "pad_width": [((256, 255), (128, 127)), ((0, 0), (6, 7))],
-        "mode": CudaPadding.supported_modes if __has_pycuda__ else [],
-        "constant_values": [0, ((1.0, 2.0), (3.0, 4.0))],
-        "output_is_none": [True, False],
-    }
-)
-
-
-@pytest.fixture(scope="class")
-def bootstrap(request):
-    cls = request.cls
-    cls.data = get_data("brain_phantom.npz")["data"]
-    cls.tol = 1e-7
-    cls.ctx = get_cuda_context(cleanup_at_exit=False)
-    yield
-    cls.ctx.pop()
-
-
-@pytest.mark.skipif(not (__has_pycuda__), reason="Need Cuda and pycuda for this test")
-@pytest.mark.usefixtures("bootstrap")
-class TestCudaPadding:
-    @pytest.mark.parametrize("config", scenarios)
-    def test_padding(self, config):
-        shape = config["shape"]
-        data = self.data[: shape[0], : shape[1]]
-        kwargs = {}
-        if config["mode"] == "constant":
-            kwargs["constant_values"] = config["constant_values"]
-        ref = np.pad(data, config["pad_width"], mode=config["mode"], **kwargs)
-        if config["output_is_none"]:
-            output = None
-        else:
-            output = garray.zeros(ref.shape, "f")
-        cuda_padding = CudaPadding(
-            config["shape"],
-            config["pad_width"],
-            mode=config["mode"],
-            constant_values=config["constant_values"],
-            cuda_options={"ctx": self.ctx},
-        )
-        d_img = garray.to_gpu(np.ascontiguousarray(data, dtype="f"))
-        res = cuda_padding.pad(d_img, output=output)
-
-        err_max = np.max(np.abs(res.get() - ref))
-        assert err_max < self.tol, str("Something wrong with padding for configuration %s" % (str(config)))
-
-    def test_custom_coordinate_transform(self):
-        data = self.data
-        R, C = np.indices(data.shape, dtype=np.int32)
-
-        pad_width = ((256, 255), (254, 251))
-        mode = "reflect"
-
-        coords_R = np.pad(R, pad_width, mode=mode)
-        coords_C = np.pad(C, pad_width, mode=mode)
-        # Further transform of coordinates - here FFT layout
-        coords_R = np.roll(coords_R, (-pad_width[0][0], -pad_width[1][0]), axis=(0, 1))
-        coords_C = np.roll(coords_C, (-pad_width[0][0], -pad_width[1][0]), axis=(0, 1))
-
-        cuda_padding = CudaPadding(data.shape, (coords_R, coords_C), mode=mode, cuda_options={"ctx": self.ctx})
-        d_img = garray.to_gpu(data)
-        d_out = garray.zeros(cuda_padding.padded_shape, "f")
-        res = cuda_padding.pad(d_img, output=d_out)
-
-        ref = np.roll(np.pad(data, pad_width, mode=mode), (-pad_width[0][0], -pad_width[1][0]), axis=(0, 1))
-
-        err_max = np.max(np.abs(d_out.get() - ref))
-        assert err_max < self.tol, "Something wrong with custom padding"

nabu/processing/tests/test_roll.py

@@ -0,0 +1,63 @@
+import numpy as np
+import pytest
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+from nabu.opencl.utils import __has_pyopencl__, get_opencl_context
+from nabu.testutils import get_data, generate_tests_scenarios, __do_long_tests__
+from nabu.processing.roll_opencl import OpenCLRoll
+
+configs_roll = {
+    "shape": [(300, 451), (300, 300), (255, 300)],
+    "offset_x": [0, 10, 155],
+    "dtype": [np.float32],  # , np.complex64],
+}
+
+
+scenarios_roll = generate_tests_scenarios(configs_roll)
+
+
+@pytest.fixture(scope="class")
+def bootstrap_roll(request):
+    cls = request.cls
+    cls.data = get_data("chelsea.npz")["data"]
+    cls.tol = 1e-7
+    if __has_pycuda__:
+        cls.cu_ctx = get_cuda_context(cleanup_at_exit=False)
+    if __has_pyopencl__:
+        cls.cl_ctx = get_opencl_context(device_type="all")
+    yield
+    if __has_pycuda__:
+        cls.cu_ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap_roll")
+class TestRoll:
+    @staticmethod
+    def _compute_ref(data, direction, offset):
+        ref = data.copy()
+        ref[:, offset:] = np.roll(data[:, offset:], direction, axis=1)
+        return ref
+
+    @pytest.mark.skipif(not (__has_pyopencl__), reason="Need pyopencl for this test")
+    @pytest.mark.parametrize("config", scenarios_roll)
+    def test_opencl_roll(self, config):
+        shape = config["shape"]
+        dtype = config["dtype"]
+        offset_x = config["offset_x"]
+        data = np.ascontiguousarray(self.data[: shape[0], : shape[1]], dtype=dtype)
+
+        ref_forward = self._compute_ref(data, 1, offset_x)
+        ref_backward = self._compute_ref(data, -1, offset_x)
+
+        roll_forward = OpenCLRoll(dtype, direction=1, offset=offset_x, ctx=self.cl_ctx)
+        d_data = roll_forward.processing.allocate_array("data", data.shape, dtype=dtype)
+        d_data.set(data)
+        roll_backward = OpenCLRoll(dtype, direction=-1, offset=offset_x, queue=roll_forward.processing.queue)
+
+        roll_forward(d_data)
+        # from spire.utils import ims
+        # ims([d_data.get(), ref_forward, d_data.get() - ref_forward])
+        assert np.allclose(d_data.get(), ref_forward), "roll_forward: something wrong with config=%s" % (str(config))
+
+        d_data.set(data)
+        roll_backward(d_data)
+        assert np.allclose(d_data.get(), ref_backward), "roll_backward: something wrong with config=%s" % (str(config))