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

build/lib/build/lib/nabu/reconstruction/rings_cuda.py

@@ -1,322 +0,0 @@
-import numpy as np
-from ..utils import docstring, get_cuda_srcfile, updiv
-from ..cuda.processing import CudaProcessing, __has_pycuda__
-from ..processing.padding_cuda import CudaPadding
-from ..processing.fft_cuda import get_fft_class, get_available_fft_implems
-from ..processing.transpose import CudaTranspose
-from ..thirdparty.tomocupy_remove_stripe import remove_all_stripe_pycuda, __have_tomocupy_deringer__
-from .rings import MunchDeringer, SinoMeanDeringer, VoDeringer
-
-if __has_pycuda__:
-    import pycuda.gpuarray as garray
-    from ..cuda.kernel import CudaKernel
-
-try:
-    from pycudwt import Wavelets
-
-    __have_pycudwt__ = True
-except ImportError:
-    __have_pycudwt__ = False
-
-
-class CudaMunchDeringer(MunchDeringer):
-    def __init__(
-        self,
-        sigma,
-        sinos_shape,
-        levels=None,
-        wname="db15",
-        padding=None,
-        padding_mode="edge",
-        fft_backend="skcuda",
-        cuda_options=None,
-    ):
-        """
-        Initialize a "Munch Et Al" sinogram deringer with the Cuda backend.
-        See References for more information.
-
-        Parameters
-        -----------
-        sigma: float
-            Standard deviation of the damping parameter. The higher value of sigma,
-            the more important the filtering effect on the rings.
-        levels: int, optional
-            Number of wavelets decomposition levels.
-            By default (None), the maximum number of decomposition levels is used.
-        wname: str, optional
-            Default is "db15" (Daubechies, 15 vanishing moments)
-        sinos_shape: tuple, optional
-            Shape of the sinogram (or sinograms stack).
-
-        References
-        ----------
-        B. Munch, P. Trtik, F. Marone, M. Stampanoni, Stripe and ring artifact removal with
-        combined wavelet-Fourier filtering, Optics Express 17(10):8567-8591, 2009.
-        """
-        super().__init__(sigma, sinos_shape, levels=levels, wname=wname, padding=padding, padding_mode=padding_mode)
-        self._check_can_use_wavelets()
-        self.cuda_processing = CudaProcessing(**(cuda_options or {}))
-        self.ctx = self.cuda_processing.ctx
-        self._init_pycudwt()
-        self._init_padding()
-        self._init_fft(fft_backend)
-        self._setup_fw_kernel()
-
-    def _check_can_use_wavelets(self):
-        if not (__have_pycudwt__):
-            raise ValueError("Needs pycudwt to use this class")
-
-    def _init_padding(self):
-        if self.padding is None:
-            return
-        self.padder = CudaPadding(
-            self.sinos_shape[1:],
-            ((0, 0), self.padding),
-            mode=self.padding_mode,
-            cuda_options={"ctx": self.cuda_processing.ctx},
-        )
-
-    def _init_fft(self, fft_backend):
-        self.fft_cls = get_fft_class(backend=fft_backend)
-        # For all k >= 1, we perform a batched (I)FFT along axis 0 on an array
-        # of shape (n_a/2^k, n_x/2^k)  (up to DWT size rounding)
-        if self.fft_cls.implem == "vkfft":
-            self._create_plans_vkfft()
-        else:
-            self._create_plans_skfft()
-
-    def _create_plans_skfft(self):
-        self._fft_plans = {}
-        for level, d_vcoeff in self._d_vertical_coeffs.items():
-            self._fft_plans[level] = self.fft_cls(d_vcoeff.shape, np.float32, r2c=True, axes=(0,), ctx=self.ctx)
-
-    def _create_plans_vkfft(self):
-        """
-        VKFFT does not support batched R2C transforms along axis 0 ("slow axis").
-        We can either use C2C (faster, but needs more memory) or transpose the arrays to do R2C along axis=1.
-        Here we transpose the arrays.
-        """
-        self._fft_plans = {}
-        self._transpose_forward_1 = {}
-        self._transpose_forward_2 = {}
-        self._transpose_inverse_1 = {}
-        self._transpose_inverse_2 = {}
-        for level, d_vcoeff in self._d_vertical_coeffs.items():
-            shape = d_vcoeff.shape
-            # Normally, a batched 1D fft on 2D data of shape (Ny, Nx) along axis 0 returns an array of shape (Ny/2+1, Nx):
-            #
-            #  (Ny, Nx)  --[fft_0]--> (Ny/2, Nx)
-            #    f32                      c64
-            #
-            # In this case, we can only do batched 1D transform along axis 1, so we have to trick with transposes:
-            #
-            #  (Ny, Nx) --[T]--> (Nx, Ny) --[fft_1]--> (Nx, Ny/2) --[T]--> (Ny/2, Nx)
-            #    f32                f32                   c64                  c64
-            #
-            # (In both cases IFFT is done the same way from right to left)
-            self._transpose_forward_1[level] = CudaTranspose(shape, np.float32, ctx=self.ctx)
-            self._fft_plans[level] = self.fft_cls(shape[::-1], np.float32, r2c=True, ctx=self.ctx)
-            self._transpose_forward_2[level] = CudaTranspose((shape[1], shape[0] // 2 + 1), np.complex64, ctx=self.ctx)
-            self._transpose_inverse_1[level] = CudaTranspose((shape[0] // 2 + 1, shape[1]), np.complex64, ctx=self.ctx)
-            self._transpose_inverse_2[level] = CudaTranspose(shape[::-1], np.float32, ctx=self.ctx)
-
-    def _init_pycudwt(self):
-        if self.levels is None:
-            self.levels = 100  # will be clipped by pycudwt
-        sino_shape = self.sinos_shape[1:] if self.padding is None else self.sino_padded_shape
-        self.cudwt = Wavelets(np.zeros(sino_shape, "f"), self.wname, self.levels)
-        self.levels = self.cudwt.levels
-        # Access memory allocated by "pypwt" from pycuda
-        self._d_sino = garray.empty(sino_shape, np.float32, gpudata=self.cudwt.image_int_ptr())
-        self._get_vertical_coeffs()
-
-    def _get_vertical_coeffs(self):
-        self._d_vertical_coeffs = {}
-        # Transfer the (0-memset) coefficients in order to get all the shapes
-        coeffs = self.cudwt.coeffs
-        for i in range(self.cudwt.levels):
-            shape = coeffs[i + 1][1].shape
-            self._d_vertical_coeffs[i + 1] = garray.empty(
-                shape, np.float32, gpudata=self.cudwt.coeff_int_ptr(3 * i + 2)
-            )
-
-    def _setup_fw_kernel(self):
-        self._fw_kernel = CudaKernel(
-            "kern_fourierwavelets",
-            filename=get_cuda_srcfile("fourier_wavelets.cu"),
-            signature="Piif",
-        )
-
-    def _apply_fft(self, level):
-        d_coeffs = self._d_vertical_coeffs[level]
-        # All the memory is allocated (or re-used) under the hood
-        if self.fft_cls.implem == "vkfft":
-            d_coeffs_t = self._transpose_forward_1[level](
-                d_coeffs
-            )  # allocates self._transpose_forward_1[level].processing.dst
-            d_coeffs_t_f = self._fft_plans[level].fft(d_coeffs_t)  # allocates self._fft_plans[level].output_fft
-            d_coeffs_f = self._transpose_forward_2[level](
-                d_coeffs_t_f
-            )  # allocates self._transpose_forward_2[level].processing.dst
-        else:
-            d_coeffs_f = self._fft_plans[level].fft(d_coeffs)
-        return d_coeffs_f
-
-    def _apply_ifft(self, d_coeffs_f, level):
-        d_coeffs = self._d_vertical_coeffs[level]
-        if self.fft_cls.implem == "vkfft":
-            d_coeffs_t_f = self._transpose_inverse_1[level](d_coeffs_f, dst=self._fft_plans[level].output_fft)
-            d_coeffs_t = self._fft_plans[level].ifft(
-                d_coeffs_t_f, output=self._transpose_forward_1[level].processing.dst
-            )
-            self._transpose_inverse_2[level](d_coeffs_t, dst=d_coeffs)
-        else:
-            self._fft_plans[level].ifft(d_coeffs_f, output=d_coeffs)
-
-    def _destripe_2D(self, d_sino, output):
-        if self.padding is not None:
-            d_sino = self.padder.pad(d_sino)
-        # set the "image" for DWT (memcpy D2D)
-        self._d_sino.set(d_sino)
-        # perform forward DWT
-        self.cudwt.forward()
-        for i in range(self.cudwt.levels):
-            level = i + 1
-            Ny, Nx = self._d_vertical_coeffs[level].shape
-            # Batched FFT along axis 0
-            d_vertical_coeffs_f = self._apply_fft(level)
-
-            # Dampen the wavelets coefficients
-            self._fw_kernel(d_vertical_coeffs_f, Nx, Ny, self.sigma)
-            # IFFT
-            self._apply_ifft(d_vertical_coeffs_f, level)
-
-        # Finally, inverse DWT
-        self.cudwt.inverse()
-        d_out = self._d_sino
-        if self.padding is not None:
-            d_out = self._d_sino[:, self.padding[0] : -self.padding[1]]  # memcpy2D
-        output.set(d_out)
-        return output
-
-
-def can_use_cuda_deringer():
-    """
-    Check wether cuda implementation of deringer can be used.
-    Checking for installed modules is not enough, as for example pyvkfft can be installed without cuda devices
-    """
-    can_do_fft = get_available_fft_implems() != []
-    return can_do_fft and __have_pycudwt__
-
-
-class CudaVoDeringer(VoDeringer):
-    """
-    An interface to topocupy's "remove_all_stripe".
-    """
-
-    def _check_requirement(self):
-        if not (__have_tomocupy_deringer__):
-            raise ImportError("need cupy")
-
-    def remove_rings_radios(self, radios):
-        return remove_all_stripe_pycuda(radios, **self._remove_all_stripe_kwargs)
-
-    def remove_rings_sinograms(self, sinos):
-        radios = sinos.transpose(axes=(1, 0, 2))  # view, no copy
-        self.remove_rings_radios(radios)
-        return sinos
-
-    def remove_rings_sinogram(self, sino):
-        radios = sino.reshape(sino.shape[0], 1, sino.shape[1])  # no copy
-        self.remove_rings_radios(radios)
-        return sino
-
-    remove_rings = remove_rings_sinograms
-
-
-class CudaSinoMeanDeringer(SinoMeanDeringer):
-    @docstring(SinoMeanDeringer)
-    def __init__(
-        self,
-        sinos_shape,
-        mode="subtract",
-        filter_cutoff=None,
-        padding_mode="edge",
-        fft_num_threads=None,
-        **cuda_options,
-    ):
-        self.processing = CudaProcessing(**(cuda_options or {}))
-        super().__init__(sinos_shape, mode, filter_cutoff, padding_mode, fft_num_threads)
-        self._init_kernels()
-
-    def _init_kernels(self):
-        self.d_sino_profile = self.processing.allocate_array("sino_profile", self.n_x)
-        self._mean_kernel = self.processing.kernel(
-            "vertical_mean",
-            filename=get_cuda_srcfile("normalization.cu"),
-            signature="PPiii",
-        )
-        self._mean_kernel_block = (32, 1, 32)
-        self._mean_kernel_grid = [updiv(a, b) for a, b in zip(self.sinos_shape[::-1], self._mean_kernel_block)]
-        self._mean_kernel_args = [self.d_sino_profile, np.int32(self.n_x), np.int32(self.n_angles), np.int32(self.n_z)]
-        self._mean_kernel_kwargs = {
-            "grid": self._mean_kernel_grid,
-            "block": self._mean_kernel_block,
-        }
-
-        self._op_kernel = self.processing.kernel(
-            "inplace_generic_op_3Dby1D",
-            filename=get_cuda_srcfile("ElementOp.cu"),
-            signature="PPiii",
-            options=["-DGENERIC_OP=%d" % (3 if self.mode == "divide" else 1)],
-        )
-        self._op_kernel_block = (16, 16, 4)
-        self._op_kernel_grid = [updiv(a, b) for a, b in zip(self.sinos_shape[::-1], self._op_kernel_block)]
-        self._op_kernel_args = [self.d_sino_profile, np.int32(self.n_x), np.int32(self.n_angles), np.int32(self.n_z)]
-        self._op_kernel_kwargs = {
-            "grid": self._op_kernel_grid,
-            "block": self._op_kernel_block,
-        }
-
-    def _init_filter(self, filter_cutoff, fft_num_threads, padding_mode):
-        super()._init_filter(filter_cutoff, fft_num_threads, padding_mode)
-        if filter_cutoff is None:
-            return
-        self._d_filter_f = self.processing.to_device("_filter_f", self._filter_f)
-
-        self.padder = CudaPadding(
-            (self.n_x, 1),
-            ((self._pad_left, self._pad_right), (0, 0)),
-            mode=self.padding_mode,
-            cuda_options={"ctx": self.processing.ctx},
-        )
-        fft_cls = get_fft_class()
-        self._fft = fft_cls(self._filter_size, np.float32, r2c=True)
-
-    def _apply_filter(self, sino_profile):
-        if self._filter_f is None:
-            return sino_profile
-
-        sino_profile = sino_profile.reshape((-1, 1))  # view
-        sino_profile_p = self.padder.pad(sino_profile).ravel()
-
-        sino_profile_f = self._fft.fft(sino_profile_p)
-        sino_profile_f *= self._d_filter_f
-        self._fft.ifft(sino_profile_f, output=sino_profile_p)
-
-        self.d_sino_profile[:] = sino_profile_p[self._pad_left : -self._pad_right]
-        return self.d_sino_profile
-
-    def remove_rings_sinogram(self, sino, output=None):
-        #
-        if output is not None:
-            raise NotImplementedError
-        #
-        self._mean_kernel(sino, *self._mean_kernel_args, **self._mean_kernel_kwargs)
-        self._apply_filter(self.d_sino_profile)
-        self._op_kernel(sino, *self._op_kernel_args, **self._op_kernel_kwargs)
-        return sino
-
-    def remove_rings_sinograms(self, sinograms):
-        for i in range(sinograms.shape[0]):
-            self.remove_rings_sinogram(sinograms[i])