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

build/lib/build/lib/nabu/preproc/tests/test_vshift.py

@@ -0,0 +1,73 @@
+import pytest
+import numpy as np
+from scipy.ndimage import shift as ndshift
+from nabu.preproc.shift import VerticalShift
+from nabu.cuda.utils import __has_pycuda__, get_cuda_context
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+    from nabu.preproc.shift_cuda import CudaVerticalShift
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    data = np.zeros([13, 11], "f")
+    slope = 100 + np.arange(13)
+    data[:] = slope[:, None]
+    cls.radios = np.array([data] * 17)
+    cls.shifts = 0.3 + np.arange(17)
+    cls.indexes = range(17)
+    # given the shifts and the radios we build the golden reference
+    golden = []
+    for iradio in range(17):
+        projection_number = cls.indexes[iradio]
+        my_shift = cls.shifts[projection_number]
+        padded_radio = np.concatenate(
+            [cls.radios[iradio], np.zeros([1, 11], "f")], axis=0
+        )  # needs padding because ndshifs does not work as expected
+        shifted_padded_radio = ndshift(padded_radio, [-my_shift, 0], mode="constant", cval=0.0, order=1).astype("f")
+        shifted_radio = shifted_padded_radio[:-1]
+        golden.append(shifted_radio)
+    cls.golden = np.array(golden)
+    cls.tol = 1e-5
+    if __has_pycuda__:
+        cls.ctx = get_cuda_context()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestVerticalShift:
+    def test_vshift(self):
+        radios = self.radios.copy()
+        new_radios = np.zeros_like(radios)
+
+        Shifter = VerticalShift(radios.shape, self.shifts)
+
+        Shifter.apply_vertical_shifts(radios, self.indexes, output=new_radios)
+        assert abs(new_radios - self.golden).max() < self.tol
+
+        Shifter.apply_vertical_shifts(radios, self.indexes)
+        assert abs(radios - self.golden).max() < self.tol
+
+    @pytest.mark.skipif(not (__has_pycuda__), reason="Need cuda/pycuda for this test")
+    def test_cuda_vshift(self):
+        d_radios = garray.to_gpu(self.radios)
+        d_radios2 = d_radios.copy()
+        d_out = garray.zeros_like(d_radios)
+
+        Shifter = CudaVerticalShift(d_radios.shape, self.shifts)
+        Shifter.apply_vertical_shifts(d_radios, self.indexes, output=d_out)
+        assert abs(d_out.get() - self.golden).max() < self.tol
+
+        Shifter.apply_vertical_shifts(d_radios, self.indexes)
+        assert abs(d_radios.get() - self.golden).max() < self.tol
+
+        # Test with negative shifts
+        radios2 = self.radios.copy()
+        Shifter_neg = VerticalShift(self.radios.shape, -self.shifts)
+        Shifter_neg.apply_vertical_shifts(radios2, self.indexes)
+
+        Shifter_neg_cuda = CudaVerticalShift(d_radios.shape, -self.shifts)
+        Shifter_neg_cuda.apply_vertical_shifts(d_radios2, self.indexes)
+        err_max = np.max(np.abs(d_radios2.get() - radios2))
+        assert err_max < 1e-6, "Something wrong for negative translations: max error = %.2e" % err_max

build/lib/build/lib/nabu/processing/convolution_cuda.py

@@ -0,0 +1,375 @@
+from os.path import dirname
+import numpy as np
+from ..utils import updiv, get_cuda_srcfile
+from ..cuda.utils import __has_pycuda__
+from ..misc.utils import ConvolutionInfos
+from ..cuda.processing import CudaProcessing
+
+if __has_pycuda__:
+    from pycuda.compiler import SourceModule
+
+
+class Convolution:
+    """
+    A class for performing convolution on GPU with CUDA, but without using
+    textures (unlike for example in ``silx.opencl.convolution``)
+    """
+
+    def __init__(self, shape, kernel, axes=None, mode=None, extra_options=None, cuda_options=None):
+        """
+        Constructor of Cuda Convolution.
+
+        Parameters
+        -----------
+        shape: tuple
+            Shape of the array.
+        kernel: array-like
+            Convolution kernel (1D, 2D or 3D).
+        axes: tuple, optional
+            Axes along which the convolution is performed,
+            for batched convolutions.
+        mode: str, optional
+            Boundary handling mode. Available modes are:
+               - "reflect": cba|abcd|dcb
+               - "nearest": aaa|abcd|ddd
+               - "wrap": bcd|abcd|abc
+               - "constant": 000|abcd|000
+
+            Default is "reflect".
+        extra_options: dict, optional
+            Advanced options (dict). Current options are:
+               - "allocate_input_array": True
+               - "allocate_output_array": True
+               - "allocate_tmp_array": True
+               - "sourcemodule_kwargs": {}
+               - "batch_along_flat_dims": True
+        """
+        self.cuda = CudaProcessing(**(cuda_options or {}))
+        self._configure_extra_options(extra_options)
+        self._determine_use_case(shape, kernel, axes)
+        self._allocate_memory(mode)
+        self._init_kernels()
+
+    def _configure_extra_options(self, extra_options):
+        self.extra_options = {
+            "allocate_input_array": True,
+            "allocate_output_array": True,
+            "allocate_tmp_array": True,
+            "sourcemodule_kwargs": {},
+            "batch_along_flat_dims": True,
+        }
+        extra_opts = extra_options or {}
+        self.extra_options.update(extra_opts)
+        self.sourcemodule_kwargs = self.extra_options["sourcemodule_kwargs"]
+
+    def _get_dimensions(self, shape, kernel):
+        self.shape = shape
+        self.data_ndim = self._check_dimensions(shape=shape, name="Data")
+        self.kernel_ndim = self._check_dimensions(arr=kernel, name="Kernel")
+        Nx = shape[-1]
+        if self.data_ndim >= 2:
+            Ny = shape[-2]
+        else:
+            Ny = 1
+        if self.data_ndim >= 3:
+            Nz = shape[-3]
+        else:
+            Nz = 1
+        self.Nx = np.int32(Nx)
+        self.Ny = np.int32(Ny)
+        self.Nz = np.int32(Nz)
+
+    def _determine_use_case(self, shape, kernel, axes):
+        """
+        Determine the convolution use case from the input/kernel shape, and axes.
+        """
+        self._get_dimensions(shape, kernel)
+        if self.kernel_ndim > self.data_ndim:
+            raise ValueError("Kernel dimensions cannot exceed data dimensions")
+        data_ndim = self.data_ndim
+        kernel_ndim = self.kernel_ndim
+        self.kernel = kernel.astype("f")
+
+        convol_infos = ConvolutionInfos()
+        k = (data_ndim, kernel_ndim)
+        if k not in convol_infos.use_cases:
+            raise ValueError(
+                "Cannot find a use case for data ndim = %d and kernel ndim = %d" % (data_ndim, kernel_ndim)
+            )
+        possible_use_cases = convol_infos.use_cases[k]
+
+        # If some dimensions are "flat", make a batched convolution along them
+        # Ex. data_dim = (1, Nx) -> batched 1D convolution
+        if self.extra_options["batch_along_flat_dims"] and (1 in self.shape):
+            axes = tuple([curr_dim for numels, curr_dim in zip(self.shape, range(len(self.shape))) if numels != 1])
+        #
+        self.use_case_name = None
+        for uc_name, uc_params in possible_use_cases.items():
+            if axes in convol_infos.allowed_axes[uc_name]:
+                self.use_case_name = uc_name
+                self.use_case_desc = uc_params["name"]
+                self.use_case_kernels = uc_params["kernels"].copy()
+        if self.use_case_name is None:
+            raise ValueError(
+                "Cannot find a use case for data ndim = %d, kernel ndim = %d and axes=%s"
+                % (data_ndim, kernel_ndim, str(axes))
+            )
+        # TODO implement this use case
+        if self.use_case_name == "batched_separable_2D_1D_3D":
+            raise NotImplementedError("The use case %s is not implemented" % self.use_case_name)
+        #
+        self.axes = axes
+        # Replace "axes=None" with an actual value (except for ND-ND)
+        allowed_axes = convol_infos.allowed_axes[self.use_case_name]
+        if len(allowed_axes) > 1:
+            # The default choice might impact perfs
+            self.axes = allowed_axes[0] or allowed_axes[1]
+        self.separable = self.use_case_name.startswith("separable")
+        self.batched = self.use_case_name.startswith("batched")
+
+    def _allocate_memory(self, mode):
+        self.mode = mode or "reflect"
+        # The current implementation does not support kernel size bigger than data size,
+        # except for mode="nearest"
+        for i, dim_size in enumerate(self.shape):
+            if min(self.kernel.shape) > dim_size and i in self.axes:
+                print(
+                    "Warning: kernel support is too large for data dimension %d (%d). Forcing convolution mode to 'nearest'"
+                    % (i, dim_size)
+                )
+                self.mode = "nearest"
+        #
+        option_array_names = {
+            "allocate_input_array": "data_in",
+            "allocate_output_array": "data_out",
+            "allocate_tmp_array": "data_tmp",
+        }
+        # Nonseparable transforms do not need tmp array
+        if not (self.separable):
+            self.extra_options["allocate_tmp_array"] = False
+        # Allocate arrays
+        for option_name, array_name in option_array_names.items():
+            if self.extra_options[option_name]:
+                value = self.cuda.allocate_array("value", self.shape, np.float32)
+            else:
+                value = None
+            setattr(self, array_name, value)
+
+        if isinstance(self.kernel, np.ndarray):
+            self.d_kernel = self.cuda.to_device("d_kernel", self.kernel)
+        else:
+            if not (isinstance(self.kernel, self.cuda.array_class)):
+                raise ValueError("kernel must be either numpy array or pycuda array")
+            self.d_kernel = self.kernel
+        self._old_input_ref = None
+        self._old_output_ref = None
+        self._c_modes_mapping = {
+            "periodic": 2,
+            "wrap": 2,
+            "nearest": 1,
+            "replicate": 1,
+            "reflect": 0,
+            "constant": 3,
+        }
+        mp = self._c_modes_mapping
+        if self.mode.lower() not in mp:
+            raise ValueError(
+                """
+                Mode %s is not available. Available modes are:
+                %s
+                """
+                % (self.mode, str(mp.keys()))
+            )
+        if self.mode.lower() == "constant":
+            raise NotImplementedError("mode='constant' is not implemented yet")
+        self._c_conv_mode = mp[self.mode]
+
+    def _init_kernels(self):
+        if self.kernel_ndim > 1:
+            if np.abs(np.diff(self.kernel.shape)).max() > 0:
+                raise NotImplementedError("Non-separable convolution with non-square kernels is not implemented yet")
+        # Compile source module
+        compile_options = [str("-DUSED_CONV_MODE=%d" % self._c_conv_mode)]
+        fname = get_cuda_srcfile("convolution.cu")
+        nabu_cuda_dir = dirname(fname)
+        include_dirs = [nabu_cuda_dir]
+        self.sourcemodule_kwargs["options"] = compile_options
+        self.sourcemodule_kwargs["include_dirs"] = include_dirs
+        with open(fname) as fid:
+            cuda_src = fid.read()
+        self._module = SourceModule(cuda_src, **self.sourcemodule_kwargs)
+        # Blocks, grid
+        self._block_size = {1: (32, 1, 1), 2: (32, 32, 1), 3: (16, 8, 8)}[self.data_ndim]  # TODO tune
+        self._n_blocks = tuple([int(updiv(a, b)) for a, b in zip(self.shape[::-1], self._block_size)])
+        # Prepare cuda kernel calls
+        self._cudakernel_signature = {
+            1: "PPPiiii",
+            2: "PPPiiiii",
+            3: "PPPiiiiii",
+        }[self.kernel_ndim]
+        self.cuda_kernels = {}
+        for axis, kern_name in enumerate(self.use_case_kernels):
+            self.cuda_kernels[axis] = self._module.get_function(kern_name)
+            self.cuda_kernels[axis].prepare(self._cudakernel_signature)
+
+        # Cuda kernel arguments
+        kernel_args = [
+            self._n_blocks,
+            self._block_size,
+            None,
+            None,
+            self.d_kernel.gpudata,
+            np.int32(self.kernel.shape[0]),
+            self.Nx,
+            self.Ny,
+            self.Nz,
+        ]
+        if self.kernel_ndim == 2:
+            kernel_args.insert(5, np.int32(self.kernel.shape[1]))
+        if self.kernel_ndim == 3:
+            kernel_args.insert(5, np.int32(self.kernel.shape[2]))
+            kernel_args.insert(6, np.int32(self.kernel.shape[1]))
+        self.kernel_args = tuple(kernel_args)
+        # If self.data_tmp is allocated, separable transforms can be performed
+        # by a series of batched transforms, without any copy, by swapping refs.
+        self.swap_pattern = None
+        if self.separable:
+            if self.data_tmp is not None:
+                self.swap_pattern = {
+                    2: [("data_in", "data_tmp"), ("data_tmp", "data_out")],
+                    3: [
+                        ("data_in", "data_out"),
+                        ("data_out", "data_tmp"),
+                        ("data_tmp", "data_out"),
+                    ],
+                }
+            else:
+                raise NotImplementedError("For now, data_tmp has to be allocated")
+
+    def _get_swapped_arrays(self, i):
+        """
+        Get the input and output arrays to use when using a "swap pattern".
+        Swapping refs enables to avoid copies between temp. array and output.
+        For example, a separable 2D->1D convolution on 2D data reads:
+          data_tmp = convol(data_input, kernel, axis=1) # step i=0
+          data_out = convol(data_tmp, kernel, axis=0) # step i=1
+
+        :param i: current step number of the separable convolution
+        """
+        n_batchs = len(self.axes)
+        in_ref, out_ref = self.swap_pattern[n_batchs][i]
+        d_in = getattr(self, in_ref)
+        d_out = getattr(self, out_ref)
+        return d_in, d_out
+
+    def _configure_kernel_args(self, cuda_kernel_args, input_ref, output_ref):
+        # TODO more elegant
+        if isinstance(input_ref, self.cuda.array_class):
+            input_ref = input_ref.gpudata
+        if isinstance(output_ref, self.cuda.array_class):
+            output_ref = output_ref.gpudata
+        if input_ref is not None or output_ref is not None:
+            cuda_kernel_args = list(cuda_kernel_args)
+            if input_ref is not None:
+                cuda_kernel_args[2] = input_ref
+            if output_ref is not None:
+                cuda_kernel_args[3] = output_ref
+            cuda_kernel_args = tuple(cuda_kernel_args)
+        return cuda_kernel_args
+
+    @staticmethod
+    def _check_dimensions(arr=None, shape=None, name="", dim_min=1, dim_max=3):
+        if shape is not None:
+            ndim = len(shape)
+        elif arr is not None:
+            ndim = arr.ndim
+        else:
+            raise ValueError("Please provide either arr= or shape=")
+        if ndim < dim_min or ndim > dim_max:
+            raise ValueError("%s dimensions should be between %d and %d" % (name, dim_min, dim_max))
+        return ndim
+
+    def _check_array(self, arr):
+        if not (isinstance(arr, self.cuda.array_class) or isinstance(arr, np.ndarray)):
+            raise TypeError("Expected either pycuda.gpuarray or numpy.ndarray")
+        if arr.dtype != np.float32:
+            raise TypeError("Data must be float32")
+        if arr.shape != self.shape:
+            raise ValueError("Expected data shape = %s" % str(self.shape))
+
+    def _set_arrays(self, array, output=None):
+        # Either copy H->D or update references.
+        if isinstance(array, np.ndarray):
+            self.data_in[:] = array[:]
+        else:
+            self._old_input_ref = self.data_in
+            self.data_in = array
+        data_in_ref = self.data_in
+        if output is not None:
+            if not (isinstance(output, np.ndarray)):
+                self._old_output_ref = self.data_out
+                self.data_out = output
+        # Update Cuda kernel arguments with new array references
+        self.kernel_args = self._configure_kernel_args(self.kernel_args, data_in_ref, self.data_out)
+
+    def _separable_convolution(self):
+        assert len(self.axes) == len(self.use_case_kernels)
+        # Separable: one kernel call per data dimension
+        for i, axis in enumerate(self.axes):
+            in_ref, out_ref = self._get_swapped_arrays(i)
+            self._batched_convolution(axis, input_ref=in_ref, output_ref=out_ref)
+
+    def _batched_convolution(self, axis, input_ref=None, output_ref=None):
+        # Batched: one kernel call in total
+        cuda_kernel = self.cuda_kernels[axis]
+        cuda_kernel_args = self._configure_kernel_args(self.kernel_args, input_ref, output_ref)
+        ev = cuda_kernel.prepared_call(*cuda_kernel_args)
+
+    def _nd_convolution(self):
+        assert len(self.use_case_kernels) == 1
+        cuda_kernel = self._module.get_function(self.use_case_kernels[0])
+        ev = cuda_kernel.prepared_call(*self.kernel_args)
+
+    def _recover_arrays_references(self):
+        if self._old_input_ref is not None:
+            self.data_in = self._old_input_ref
+            self._old_input_ref = None
+        if self._old_output_ref is not None:
+            self.data_out = self._old_output_ref
+            self._old_output_ref = None
+        self.kernel_args = self._configure_kernel_args(self.kernel_args, self.data_in, self.data_out)
+
+    def _get_output(self, output):
+        if output is None:
+            res = self.data_out.get()
+        else:
+            res = output
+            if isinstance(output, np.ndarray):
+                output[:] = self.data_out[:]
+        self._recover_arrays_references()
+        return res
+
+    def convolve(self, array, output=None):
+        """
+        Convolve an array with the class kernel.
+
+        :param array: Input array. Can be numpy.ndarray or pycuda.gpuarray.GPUArray.
+        :param output: Output array. Can be numpy.ndarray or pycuda.gpuarray.GPUArray.
+        """
+        self._check_array(array)
+        self._set_arrays(array, output=output)
+        if self.axes is not None:
+            if self.separable:
+                self._separable_convolution()
+            elif self.batched:
+                assert len(self.axes) == 1
+                self._batched_convolution(self.axes[0])
+            # else: ND-ND convol
+        else:
+            # ND-ND convol
+            self._nd_convolution()
+
+        res = self._get_output(output)
+        return res
+
+    __call__ = convolve

build/lib/build/lib/nabu/processing/fft_base.py

@@ -0,0 +1,163 @@
+import numpy as np
+from ..utils import BaseClassError
+
+
+class _BaseFFT:
+    """
+    A base class for FFTs.
+    """
+
+    implem = "none"
+    ProcessingCls = BaseClassError
+
+    def __init__(self, shape, dtype, r2c=True, axes=None, normalize="rescale", **backend_options):
+        """
+        Base class for Fast Fourier Transform (FFT).
+
+        Parameters
+        ----------
+        shape: list of int
+            Shape of the input data
+        dtype: str or numpy.dtype
+            Data type of the input data
+        r2c: bool, optional
+            Whether to use real-to-complex transform for real-valued input. Default is True.
+        axes: list of int, optional
+            Axes along which FFT is computed.
+              * For 2D transform: axes=(1,0)
+              * For batched 1D transform of 2D image: axes=(-1,)
+        normalize: str, optional
+            Whether to normalize FFT and IFFT. Possible values are:
+              * "rescale": in this case, Fourier data is divided by "N"
+                before IFFT, so that IFFT(FFT(data)) = data.
+                This corresponds to numpy norm=None i.e norm="backward".
+              * "ortho": in this case, FFT and IFFT are adjoint of eachother,
+                the transform is unitary. Both FFT and IFFT are scaled with 1/sqrt(N).
+              * "none": no normalizatio is done : IFFT(FFT(data)) = data*N
+
+        Other parameters
+        -----------------
+        backend_options: dict, optional
+            Parameters to pass to CudaProcessing or OpenCLProcessing class.
+        """
+        self._init_backend(backend_options)
+        self._set_dtypes(dtype, r2c)
+        self._set_shape_and_axes(shape, axes)
+        self._configure_batched_transform()
+        self._configure_normalization(normalize)
+        self._compute_fft_plans()
+
+    def _init_backend(self, backend_options):
+        self.processing = self.ProcessingCls(**backend_options)
+
+    def _set_dtypes(self, dtype, r2c):
+        self.dtype = np.dtype(dtype)
+        dtypes_mapping = {
+            np.dtype("float32"): np.complex64,
+            np.dtype("float64"): np.complex128,
+            np.dtype("complex64"): np.complex64,
+            np.dtype("complex128"): np.complex128,
+        }
+        if self.dtype not in dtypes_mapping:
+            raise ValueError("Invalid input data type: got %s" % self.dtype)
+        self.dtype_out = dtypes_mapping[self.dtype]
+        self.r2c = r2c
+
+    def _set_shape_and_axes(self, shape, axes):
+        # Input shape
+        if np.isscalar(shape):
+            shape = (shape,)
+        self.shape = shape
+        # Axes
+        default_axes = tuple(range(len(self.shape)))
+        if axes is None:
+            self.axes = default_axes
+        else:
+            self.axes = tuple(np.array(default_axes)[np.array(axes)])
+        # Output shape
+        shape_out = self.shape
+        if self.r2c:
+            reduced_dim = self.axes[-1] if self.axes is not None else -1
+            shape_out = list(shape_out)
+            shape_out[reduced_dim] = shape_out[reduced_dim] // 2 + 1
+            shape_out = tuple(shape_out)
+        self.shape_out = shape_out
+
+    def _configure_batched_transform(self):
+        pass
+
+    def _configure_normalization(self, normalize):
+        pass
+
+    def _compute_fft_plans(self):
+        pass
+
+
+class _BaseVKFFT(_BaseFFT):
+    """
+    FFT using VKFFT backend
+    """
+
+    implem = "vkfft"
+    backend = "none"
+    ProcessingCls = BaseClassError
+    vkffs_cls = BaseClassError
+
+    def _configure_batched_transform(self):
+        if self.axes is not None and len(self.shape) == len(self.axes):
+            self.axes = None
+            return
+        if self.r2c:
+            # batched Real-to-complex transforms are supported only along fast axes
+            if not (is_fast_axes(len(self.shape), self.axes)):
+                raise ValueError("For %dD R2C, only batched transforms along fast axes are allowed" % (len(self.shape)))
+            self._vkfft_ndim = len(self.axes)
+            self.axes = None  # vkfft still can do a batched transform by providing dim=XX, axes=None
+
+    def _configure_normalization(self, normalize):
+        self.normalize = normalize
+        self._vkfft_norm = {
+            "rescale": 1,
+            "backward": 1,
+            "ortho": "ortho",
+            "none": 0,
+        }.get(self.normalize, 1)
+
+    def _set_shape_and_axes(self, shape, axes):
+        super()._set_shape_and_axes(shape, axes)
+        self._vkfft_ndim = None
+
+    def _compute_fft_plans(self):
+        self._vkfft_plan = self.vkffs_cls(
+            self.shape,
+            self.dtype,
+            ndim=self._vkfft_ndim,
+            inplace=False,
+            norm=self._vkfft_norm,
+            r2c=self.r2c,
+            dct=False,
+            axes=self.axes,
+            strides=None,
+            **self._vkfft_other_init_kwargs,
+        )
+
+    def fft(self, array, output=None):
+        if output is None:
+            output = self.output_fft = self.processing.allocate_array(
+                "output_fft", self.shape_out, dtype=self.dtype_out
+            )
+        return self._vkfft_plan.fft(array, dest=output)
+
+    def ifft(self, array, output=None):
+        if output is None:
+            output = self.output_ifft = self.processing.allocate_array("output_ifft", self.shape, dtype=self.dtype)
+        return self._vkfft_plan.ifft(array, dest=output)
+
+
+def is_fast_axes(ndim, axes):
+    """
+    Return true if "axes" are the fast dimensions
+    """
+    all_axes = list(range(ndim))
+    axes = sorted([ax + ndim if ax < 0 else ax for ax in axes])  # transform "-1" to an actual axis index (1 for 2D)
+    return all_axes[-len(axes) :] == axes