nabu 2024.1.0rc4__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/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

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/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)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/tests/test_histogram.py

@@ -0,0 +1,55 @@
+import pytest
+import numpy as np
+from nabu.testutils import get_data
+from nabu.processing.histogram import PartialHistogram
+from nabu.cuda.utils import __has_pycuda__, get_cuda_context
+
+if __has_pycuda__:
+    from nabu.processing.histogram_cuda import CudaPartialHistogram
+    import pycuda.gpuarray as garray
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.data = get_data("mri_rec_astra.npz")["data"]
+    cls.data /= 10
+    cls.data[:100] *= 10
+    cls.data_part_1 = cls.data[:100]
+    cls.data_part_2 = cls.data[100:]
+    cls.data0 = cls.data.ravel()
+    cls.bin_tol = 1e-5 * (cls.data0.max() - cls.data0.min())
+    cls.hist_rtol = 1.5e-3
+    if __has_pycuda__:
+        cls.ctx = get_cuda_context(cleanup_at_exit=False)
+    yield
+    if __has_pycuda__:
+        cls.ctx.pop()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestPartialHistogram:
+    def compare_histograms(self, hist1, hist2):
+        errmax_bins = np.max(np.abs(hist1[1] - hist2[1]))
+        assert errmax_bins < self.bin_tol
+        errmax_hist = np.max(np.abs(hist1[0] - hist2[0]) / hist2[0].max())
+        assert errmax_hist / hist2[0].max() < self.hist_rtol
+
+    def test_fixed_nbins(self):
+        partial_hist = PartialHistogram(method="fixed_bins_number", num_bins=1e6)
+        hist1 = partial_hist.compute_histogram(self.data_part_1.ravel())
+        hist2 = partial_hist.compute_histogram(self.data_part_2.ravel())
+        hist = partial_hist.merge_histograms([hist1, hist2])
+        ref = np.histogram(self.data0, bins=partial_hist.num_bins)
+        self.compare_histograms(hist, ref)
+
+    @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
+    def test_fixed_nbins_cuda(self):
+        partial_hist = CudaPartialHistogram(method="fixed_bins_number", num_bins=1e6, cuda_options={"ctx": self.ctx})
+        data_part_1 = garray.to_gpu(np.tile(self.data_part_1, (1, 1, 1)))
+        data_part_2 = garray.to_gpu(np.tile(self.data_part_2, (1, 1, 1)))
+        hist1 = partial_hist.compute_histogram(data_part_1)
+        hist2 = partial_hist.compute_histogram(data_part_2)
+        hist = partial_hist.merge_histograms([hist1, hist2])
+        ref = np.histogram(self.data0, bins=partial_hist.num_bins)
+        self.compare_histograms(hist, ref)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/tests/test_medfilt.py

@@ -0,0 +1,77 @@
+import pytest
+import numpy as np
+from silx.math.medianfilter import medfilt2d
+from nabu.testutils import generate_tests_scenarios, get_data
+from nabu.cuda.utils import get_cuda_context, __has_pycuda__
+
+if __has_pycuda__:
+    from nabu.processing.medfilt_cuda import MedianFilter
+    import pycuda.gpuarray as garray
+
+
+scenarios = generate_tests_scenarios(
+    {
+        "input_on_gpu": [False, True],
+        "output_on_gpu": [False, True],
+        "footprint": [(3, 3), (5, 5)],
+        "mode": ["reflect", "nearest"],
+        "batched_2d": [False, True],
+    }
+)
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls.data = np.ascontiguousarray(get_data("brain_phantom.npz")["data"][::2, ::2][:-1, :])
+    cls.tol = 1e-7
+    cls.ctx = get_cuda_context(cleanup_at_exit=False)
+    cls.allocate_numpy_arrays()
+    cls.allocate_cuda_arrays()
+    yield
+    cls.ctx.pop()
+
+
+@pytest.mark.skipif(not (__has_pycuda__), reason="Need Cuda/pycuda for this test")
+@pytest.mark.usefixtures("bootstrap")
+class TestMedianFilter(object):
+    @classmethod
+    def allocate_numpy_arrays(cls):
+        shape = cls.data.shape
+        cls.input = cls.data
+        cls.input3d = np.tile(cls.input, (2, 1, 1))
+
+    @classmethod
+    def allocate_cuda_arrays(cls):
+        shape = cls.data.shape
+        cls.d_input = garray.to_gpu(cls.input)
+        cls.d_output = garray.zeros_like(cls.d_input)
+        cls.d_input3d = garray.to_gpu(cls.input3d)
+        cls.d_output3d = garray.zeros_like(cls.d_input3d)
+
+    # parametrize on a class method will use the same class, and launch this
+    # method with different scenarios.
+    @pytest.mark.parametrize("config", scenarios)
+    def testMedfilt(self, config):
+        if config["input_on_gpu"]:
+            input_data = self.d_input if not (config["batched_2d"]) else self.d_input3d
+        else:
+            input_data = self.input if not (config["batched_2d"]) else self.input3d
+        if config["output_on_gpu"]:
+            output_data = self.d_output if not (config["batched_2d"]) else self.d_output3d
+        else:
+            output_data = None
+        # Cuda median filter
+        medfilt = MedianFilter(
+            input_data.shape,
+            footprint=config["footprint"],
+            mode=config["mode"],
+            cuda_options={"ctx": self.ctx},
+        )
+        res = medfilt.medfilt2(input_data, output=output_data)
+        if config["output_on_gpu"]:
+            res = res.get()
+        # Reference (scipy)
+        ref = medfilt2d(self.input, config["footprint"][0], mode=config["mode"])
+        max_absolute_error = np.max(np.abs(res - ref))
+        assert max_absolute_error < self.tol, "Something wrong with configuration %s" % str(config)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/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
+        )