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

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

nabu/processing/fft_cuda.py

@@ -0,0 +1,256 @@
+import os
+import warnings
+from multiprocessing import get_context
+from multiprocessing.pool import Pool
+import numpy as np
+from ..utils import check_supported
+from .fft_base import _BaseFFT, _BaseVKFFT
+
+try:
+    from pyvkfft.cuda import VkFFTApp as vk_cufft
+
+    __has_vkfft__ = True
+except (ImportError, OSError):
+    __has_vkfft__ = False
+    vk_cufft = None
+from ..cuda.processing import CudaProcessing
+
+Plan = None
+cu_fft = None
+cu_ifft = None
+__has_skcuda__ = None
+
+
+def init_skcuda():
+    # This needs to be done here, because scikit-cuda creates a Cuda context at import,
+    # which can mess things up in some cases.
+    # Ugly solution to an ugly problem.
+    global __has_skcuda__, Plan, cu_fft, cu_ifft
+    try:
+        from skcuda.fft import Plan
+        from skcuda.fft import fft as cu_fft
+        from skcuda.fft import ifft as cu_ifft
+
+        __has_skcuda__ = True
+    except ImportError:
+        __has_skcuda__ = False
+
+
+class SKCUFFT(_BaseFFT):
+    implem = "skcuda"
+    backend = "cuda"
+    ProcessingCls = CudaProcessing
+
+    def _configure_batched_transform(self):
+        if __has_skcuda__ is None:
+            init_skcuda()
+        if not (__has_skcuda__):
+            raise ImportError("Please install pycuda and scikit-cuda to use the CUDA back-end")
+
+        self.cufft_batch_size = 1
+        self.cufft_shape = self.shape
+        self._cufft_plan_kwargs = {}
+        if (self.axes is not None) and (len(self.axes) < len(self.shape)):
+            # In the easiest case, the transform is computed along the fastest dimensions:
+            #  - 1D transforms of lines of 2D data
+            #  - 2D transforms of images of 3D data (stacked along slow dim)
+            #  - 1D transforms of 3D data along fastest dim
+            # Otherwise, we have to configure cuda "advanced memory layout".
+            data_ndims = len(self.shape)
+
+            if data_ndims == 2:
+                n_y, n_x = self.shape
+                along_fast_dim = self.axes[0] == 1
+                self.cufft_shape = n_x if along_fast_dim else n_y
+                self.cufft_batch_size = n_y if along_fast_dim else n_x
+                if not (along_fast_dim):
+                    # Batched vertical 1D FFT on 2D data need advanced data layout
+                    # http://docs.nvidia.com/cuda/cufft/#advanced-data-layout
+                    self._cufft_plan_kwargs = {
+                        "inembed": np.int32([0]),
+                        "istride": n_x,
+                        "idist": 1,
+                        "onembed": np.int32([0]),
+                        "ostride": n_x,
+                        "odist": 1,
+                    }
+
+            if data_ndims == 3:
+                # TODO/FIXME - the following work for C2C but not R2C ?!
+                # fast_axes = [(1, 2), (2, 1), (2,)]
+                fast_axes = [(2,)]
+                if self.axes not in fast_axes:
+                    raise NotImplementedError(
+                        "With the CUDA backend, batched transform on 3D data is only supported along fastest dimensions"
+                    )
+                self.cufft_batch_size = self.shape[0]
+                self.cufft_shape = self.shape[1:]
+                if len(self.axes) == 1:
+                    # 1D transform on 3D data: here only supported along fast dim, so batch_size is Nx*Ny
+                    self.cufft_batch_size = np.prod(self.shape[:2])
+                    self.cufft_shape = (self.shape[-1],)
+                if len(self.cufft_shape) == 1:
+                    self.cufft_shape = self.cufft_shape[0]
+
+    def _configure_normalization(self, normalize):
+        self.normalize = normalize
+        if self.normalize == "ortho":
+            # TODO
+            raise NotImplementedError("Normalization mode 'ortho' is not implemented with CUDA backend yet.")
+        self.cufft_scale_inverse = self.normalize == "rescale"
+
+    def _compute_fft_plans(self):
+        self.plan_forward = Plan(  # pylint: disable = E1102
+            self.cufft_shape,
+            self.dtype,
+            self.dtype_out,
+            batch=self.cufft_batch_size,
+            stream=self.processing.stream,
+            **self._cufft_plan_kwargs,
+            # cufft extensible plan API is only supported after 0.5.1
+            # (commit 65288d28ca0b93e1234133f8d460dc6becb65121)
+            # but there is still no official 0.5.2
+            # ~ auto_allocate=True # cufft extensible plan API
+        )
+        self.plan_inverse = Plan(  # pylint: disable = E1102
+            self.cufft_shape,  # not shape_out
+            self.dtype_out,
+            self.dtype,
+            batch=self.cufft_batch_size,
+            stream=self.processing.stream,
+            **self._cufft_plan_kwargs,
+            # cufft extensible plan API is only supported after 0.5.1
+            # (commit 65288d28ca0b93e1234133f8d460dc6becb65121)
+            # but there is still no official 0.5.2
+            # ~ auto_allocate=True
+        )
+
+    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
+            )
+        cu_fft(array, output, self.plan_forward, scale=False)  # pylint: disable = E1102
+        return 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)
+        cu_ifft(  # pylint: disable = E1102
+            array,
+            output,
+            self.plan_inverse,
+            scale=self.cufft_scale_inverse,
+        )
+        return output
+
+
+class VKCUFFT(_BaseVKFFT):
+    """
+    Cuda FFT, using VKFFT backend
+    """
+
+    implem = "vkfft"
+    backend = "cuda"
+    ProcessingCls = CudaProcessing
+    vkffs_cls = vk_cufft
+
+    def _init_backend(self, backend_options):
+        super()._init_backend(backend_options)
+        self._vkfft_other_init_kwargs = {"stream": self.processing.stream}
+
+
+def _has_vkfft(x):
+    # should be run from within a Process
+    try:
+        from nabu.processing.fft_cuda import VKCUFFT, __has_vkfft__
+
+        if not __has_vkfft__:
+            return False
+        vk = VKCUFFT((16,), "f")
+        avail = True
+    except (RuntimeError, OSError):
+        avail = False
+    return avail
+
+
+def has_vkfft(safe=True):
+    """
+    Determine whether pyvkfft is available.
+    For Cuda GPUs, vkfft relies on nvrtc which supports a narrow range of Cuda devices.
+    Unfortunately, it's not possible to determine whether vkfft is available before creating a Cuda context.
+    So we create a process (from scratch, i.e no fork), do the test within, and exit.
+    This function cannot be tested from a notebook/console, a proper entry point has to be created (if __name__ == "__main__").
+    """
+    if not safe:
+        return _has_vkfft(None)
+    ctx = get_context("spawn")
+    with Pool(1, context=ctx) as p:
+        v = p.map(_has_vkfft, [1])[0]
+    return v
+
+
+def _has_skfft(x):
+    # should be run from within a Process
+    try:
+        from nabu.processing.fft_cuda import SKCUFFT
+
+        sk = SKCUFFT((16,), "f")
+        avail = True
+    except (ImportError, RuntimeError, OSError):
+        avail = False
+    return avail
+
+
+def has_skcuda(safe=True):
+    """
+    Determine whether scikit-cuda/CUFFT is available.
+    Currently, scikit-cuda will create a Cuda context for Cublas, which can mess up the current execution.
+    Do it in a separate thread.
+    """
+    if not safe:
+        return _has_skfft(None)
+    ctx = get_context("spawn")
+    with Pool(1, context=ctx) as p:
+        v = p.map(_has_skfft, [1])[0]
+    return v
+
+
+def get_fft_class(backend="skcuda"):
+    backends = {
+        "scikit-cuda": SKCUFFT,
+        "skcuda": SKCUFFT,
+        "cufft": SKCUFFT,
+        "scikit": SKCUFFT,
+        "vkfft": VKCUFFT,
+        "pyvkfft": VKCUFFT,
+    }
+
+    def check_vkfft(asked_fft_cls):
+        if asked_fft_cls is VKCUFFT:
+            if has_vkfft(safe=True) is False:
+                warnings.warn("Could not get VKFFT backend. Falling-back to scikit-cuda/CUFFT instead.", RuntimeWarning)
+                return SKCUFFT
+            return VKCUFFT
+        return SKCUFFT
+
+    def get_fft_cls(asked_fft_backend):
+        asked_fft_backend = asked_fft_backend.lower()
+        check_supported(asked_fft_backend, list(backends.keys()), "FFT backend name")
+        asked_fft_cls = backends[asked_fft_backend]
+        fft_cls = check_vkfft(asked_fft_cls)
+        return fft_cls
+
+    asked_fft_backend_env = os.environ.get("NABU_FFT_BACKEND", "")
+    if asked_fft_backend_env != "":
+        return get_fft_cls(asked_fft_backend_env)
+    return get_fft_cls(backend)
+
+
+def get_available_fft_implems():
+    avail_implems = []
+    if has_skcuda(safe=True):
+        avail_implems.append("skcuda")
+    if has_vkfft(safe=True):
+        avail_implems.append("vkfft")
+    return avail_implems

nabu/processing/fft_opencl.py

@@ -0,0 +1,54 @@
+from multiprocessing import get_context
+from multiprocessing.pool import Pool
+from .fft_base import _BaseVKFFT
+from ..opencl.processing import OpenCLProcessing
+
+try:
+    from pyvkfft.opencl import VkFFTApp as vk_clfft
+
+    __has_vkfft__ = True
+except (ImportError, OSError):
+    __has_vkfft__ = False
+    vk_clfft = None
+
+
+class VKCLFFT(_BaseVKFFT):
+    """
+    OpenCL FFT, using VKFFT backend
+    """
+
+    implem = "vkfft"
+    backend = "opencl"
+    ProcessingCls = OpenCLProcessing
+    vkffs_cls = vk_clfft
+
+    def _init_backend(self, backend_options):
+        super()._init_backend(backend_options)
+        self._vkfft_other_init_kwargs = {"queue": self.processing.queue}
+
+
+def _has_vkfft(x):
+    # should be run from within a Process
+    try:
+        from nabu.processing.fft_opencl import VKCLFFT, __has_vkfft__
+
+        if not __has_vkfft__:
+            return False
+        vk = VKCLFFT((16,), "f")
+        avail = True
+    except (RuntimeError, OSError):
+        avail = False
+    return avail
+
+
+def has_vkfft(safe=True):
+    """
+    Determine whether pyvkfft is available.
+    This function cannot be tested from a notebook/console, a proper entry point has to be created (if __name__ == "__main__").
+    """
+    if not safe:
+        return _has_vkfft(None)
+    ctx = get_context("spawn")
+    with Pool(1, context=ctx) as p:
+        v = p.map(_has_vkfft, [1])[0]
+    return v

nabu/processing/fftshift.py

@@ -0,0 +1,134 @@
+import numpy as np
+from ..utils import BaseClassError, get_opencl_srcfile, updiv
+from ..opencl.kernel import OpenCLKernel
+from ..opencl.processing import OpenCLProcessing
+from pyopencl.tools import dtype_to_ctype as cl_dtype_to_ctype
+
+
+class FFTshiftBase:
+    KernelCls = BaseClassError
+    ProcessingCls = BaseClassError
+    dtype_to_ctype = BaseClassError
+    backend = "none"
+
+    def __init__(self, shape, dtype, dst_dtype=None, axes=None, **backend_options):
+        """
+
+        Parameters
+        ----------
+        shape: tuple
+            Array shape - can be 1D or 2D. 3D is not supported.
+        dtype: str or numpy.dtype
+            Data type, eg. "f", numpy.complex64, ...
+        dst_dtype: str or numpy.dtype
+            Output data type. If not provided (default), the shift is done in-place.
+        axes: tuple, optional
+            Axes over which to shift.  Default is None, which shifts all axes.
+
+        Other parameters
+        ----------------
+        backend_options:
+            named arguments to pass to CudaProcessing or OpenCLProcessing
+        """
+        #
+        if axes not in [1, (1,), (-1,)]:
+            raise NotImplementedError
+        #
+        self.processing = self.ProcessingCls(**backend_options)
+        self.shape = shape
+        if len(self.shape) not in [1, 2]:
+            raise ValueError("Expected 1D or 2D array")
+        self.dtype = np.dtype(dtype)
+        self.dst_dtype = dst_dtype
+
+        if dst_dtype is None:
+            self._configure_inplace_shift()
+        else:
+            self._configure_out_of_place_shift()
+        self._configure_kenel_initialization()
+        self._fftshift_kernel = self.KernelCls(*self._kernel_init_args, **self._kernel_init_kwargs)
+        self._configure_kernel_call()
+
+    def _configure_inplace_shift(self):
+        self.inplace = True
+        # in-place on odd-sized array is more difficult - see fftshift.cl
+        if self.shape[-1] & 1:
+            raise NotImplementedError
+        #
+        self._kernel_init_args = [
+            "fftshift_x_inplace",
+        ]
+        self._kernel_init_kwargs = {
+            "options": [
+                "-DDTYPE=%s" % self.dtype_to_ctype(self.dtype),
+            ],
+        }
+
+    def _configure_out_of_place_shift(self):
+        self.inplace = False
+        self._kernel_init_args = [
+            "fftshift_x",
+        ]
+        self._kernel_init_kwargs = {
+            "options": [
+                "-DDTYPE=%s" % self.dtype_to_ctype(self.dtype),
+                "-DDTYPE_OUT=%s" % self.dtype_to_ctype(np.dtype(self.dst_dtype)),
+            ],
+        }
+        additional_flag = None
+        input_is_complex = np.iscomplexobj(np.ones(1, dtype=self.dtype))
+        output_is_complex = np.iscomplexobj(np.ones(1, dtype=self.dst_dtype))
+        if not (input_is_complex) and output_is_complex:
+            additional_flag = "-DCAST_TO_COMPLEX"
+        if input_is_complex and not (output_is_complex):
+            additional_flag = "-DCAST_TO_REAL"
+        if additional_flag is not None:
+            self._kernel_init_kwargs["options"].append(additional_flag)
+
+    def _call_fftshift_inplace(self, arr, direction):
+        self._fftshift_kernel(  # pylint: disable=E1102
+            arr, np.int32(self.shape[1]), np.int32(self.shape[0]), np.int32(direction), **self._kernel_kwargs
+        )
+        return arr
+
+    def _call_fftshift_out_of_place(self, arr, dst, direction):
+        if dst is None:
+            dst = self.processing.allocate_array("dst", arr.shape, dtype=self.dst_dtype)
+        self._fftshift_kernel(  # pylint: disable=E1102
+            arr, dst, np.int32(self.shape[1]), np.int32(self.shape[0]), np.int32(direction), **self._kernel_kwargs
+        )
+        return dst
+
+    def fftshift(self, arr, dst=None):
+        if self.inplace:
+            return self._call_fftshift_inplace(arr, 1)
+        else:
+            return self._call_fftshift_out_of_place(arr, dst, 1)
+
+    def ifftshift(self, arr, dst=None):
+        if self.inplace:
+            return self._call_fftshift_inplace(arr, -1)
+        else:
+            return self._call_fftshift_out_of_place(arr, dst, -1)
+
+
+class OpenCLFFTshift(FFTshiftBase):
+    KernelCls = OpenCLKernel
+    ProcessingCls = OpenCLProcessing
+    dtype_to_ctype = cl_dtype_to_ctype
+    backend = "opencl"
+
+    def _configure_kenel_initialization(self):
+        self._kernel_init_args.append(self.processing.ctx)
+        self._kernel_init_kwargs.update(
+            {
+                "filename": get_opencl_srcfile("fftshift.cl"),
+                "queue": self.processing.queue,
+            }
+        )
+
+    def _configure_kernel_call(self):
+        # TODO in-place fftshift needs to launch only arr.size//2 threads
+        block = (16, 16, 1)
+        grid = [updiv(a, b) * b for a, b in zip(self.shape[::-1], block)]
+        self._kernel_kwargs = {"global_size": grid, "local_size": block}