nabu 2024.1.1__py3-none-any.whl → 2024.1.2__py3-none-any.whl

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

@@ -1,73 +0,0 @@
-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/build/lib/build/lib/build/lib/nabu/processing/convolution_cuda.py

@@ -1,375 +0,0 @@
-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/build/lib/build/lib/build/lib/nabu/processing/fft_base.py

@@ -1,163 +0,0 @@
-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