from numbers import Number import warnings import numpy as np import cupy from cupy.fft._fft import (_fft, _default_fft_func, hfft as _hfft, ihfft as _ihfft, _swap_direction) _scipy_150 = False _scipy_160 = False try: import scipy import scipy.fft as _scipy_fft except ImportError: class _DummyModule: def __getattr__(self, name): return None _scipy_fft = _DummyModule() else: from numpy.lib import NumpyVersion as Version _scipy_150 = Version(scipy.__version__) >= Version('1.5.0') _scipy_160 = Version(scipy.__version__) >= Version('1.6.0') del Version del scipy # Backend support for scipy.fft __ua_domain__ = 'numpy.scipy.fft' _implemented: dict = {} def __ua_convert__(dispatchables, coerce): if coerce: try: replaced = [ cupy.asarray(d.value) if d.coercible and d.type is np.ndarray else d.value for d in dispatchables] except TypeError: return NotImplemented else: replaced = [d.value for d in dispatchables] if not all(d.type is not np.ndarray or isinstance(r, cupy.ndarray) for r, d in zip(replaced, dispatchables)): return NotImplemented return replaced def __ua_function__(method, args, kwargs): fn = _implemented.get(method, None) if fn is None: return NotImplemented if 'plan' in kwargs and not _scipy_150: warnings.warn('The \'plan\' argument is supported in SciPy v1.5.0+') return fn(*args, **kwargs) def _implements(scipy_func): """Decorator adds function to the dictionary of implemented functions""" def inner(func): _implemented[scipy_func] = func return func return inner def _assequence(x): """Convert scalars to a sequence, otherwise pass through ``x`` unchanged""" if isinstance(x, Number): return (x,) return x @_implements(_scipy_fft.fft) def fft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the one-dimensional FFT. Args: x (cupy.ndarray): Array to be transformed. n (None or int): Length of the transformed axis of the output. If ``n`` is not given, the length of the input along the axis specified by ``axis`` is used. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for transforming ``x`` over ``axis``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, n, axis) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``n`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.fft` """ from cupy.cuda import cufft return _fft(x, (n,), (axis,), norm, cufft.CUFFT_FORWARD, overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.ifft) def ifft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the one-dimensional inverse FFT. Args: x (cupy.ndarray): Array to be transformed. n (None or int): Length of the transformed axis of the output. If ``n`` is not given, the length of the input along the axis specified by ``axis`` is used. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for transforming ``x`` over ``axis``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, n, axis) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``n`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.ifft` """ from cupy.cuda import cufft return _fft(x, (n,), (axis,), norm, cufft.CUFFT_INVERSE, overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.fft2) def fft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the two-dimensional FFT. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the transformed axes of the output. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.fft2` """ return fftn(x, s, axes, norm, overwrite_x, plan=plan) @_implements(_scipy_fft.ifft2) def ifft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the two-dimensional inverse FFT. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the transformed axes of the output. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.ifft2` """ return ifftn(x, s, axes, norm, overwrite_x, plan=plan) @_implements(_scipy_fft.fftn) def fftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the N-dimensional FFT. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the transformed axes of the output. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.fftn` """ from cupy.cuda import cufft s = _assequence(s) axes = _assequence(axes) func = _default_fft_func(x, s, axes) return func(x, s, axes, norm, cufft.CUFFT_FORWARD, overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.ifftn) def ifftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the N-dimensional inverse FFT. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the transformed axes of the output. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes) Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if that of the input is another. .. seealso:: :func:`scipy.fft.ifftn` """ from cupy.cuda import cufft s = _assequence(s) axes = _assequence(axes) func = _default_fft_func(x, s, axes) return func(x, s, axes, norm, cufft.CUFFT_INVERSE, overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.rfft) def rfft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the one-dimensional FFT for real input. The returned array contains the positive frequency components of the corresponding :func:`fft`, up to and including the Nyquist frequency. Args: x (cupy.ndarray): Array to be transformed. n (None or int): Length of the transformed axis of the output. If ``n`` is not given, the length of the input along the axis specified by ``axis`` is used. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for transforming ``x`` over ``axis``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, n, axis, value_type='R2C') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array. .. seealso:: :func:`scipy.fft.rfft` """ from cupy.cuda import cufft return _fft(x, (n,), (axis,), norm, cufft.CUFFT_FORWARD, 'R2C', overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.irfft) def irfft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the one-dimensional inverse FFT for real input. Args: x (cupy.ndarray): Array to be transformed. n (None or int): Length of the transformed axis of the output. If ``n`` is not given, the length of the input along the axis specified by ``axis`` is used. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for transforming ``x`` over ``axis``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, n, axis, value_type='C2R') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array. .. seealso:: :func:`scipy.fft.irfft` """ from cupy.cuda import cufft return _fft(x, (n,), (axis,), norm, cufft.CUFFT_INVERSE, 'C2R', overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.rfft2) def rfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the two-dimensional FFT for real input. Args: a (cupy.ndarray): Array to be transform. s (None or tuple of ints): Shape to use from the input. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes, value_type='R2C') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if the input is other. The length of the last axis transformed will be ``s[-1]//2+1``. .. seealso:: :func:`scipy.fft.rfft2` """ return rfftn(x, s, axes, norm, overwrite_x, plan=plan) @_implements(_scipy_fft.irfft2) def irfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the two-dimensional inverse FFT for real input. Args: a (cupy.ndarray): Array to be transform. s (None or tuple of ints): Shape of the output. If ``s`` is not given, they are determined from the lengths of the input along the axes specified by ``axes``. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes, value_type='C2R') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if the input is other. If ``s`` is not given, the length of final transformed axis of output will be `2*(m-1)` where `m` is the length of the final transformed axis of the input. .. seealso:: :func:`scipy.fft.irfft2` """ return irfftn(x, s, axes, norm, overwrite_x, plan=plan) @_implements(_scipy_fft.rfftn) def rfftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the N-dimensional FFT for real input. Args: a (cupy.ndarray): Array to be transform. s (None or tuple of ints): Shape to use from the input. If ``s`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes, value_type='R2C') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if the input is other. The length of the last axis transformed will be ``s[-1]//2+1``. .. seealso:: :func:`scipy.fft.rfftn` """ from cupy.cuda import cufft s = _assequence(s) axes = _assequence(axes) func = _default_fft_func(x, s, axes, value_type='R2C') return func(x, s, axes, norm, cufft.CUFFT_FORWARD, 'R2C', overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.irfftn) def irfftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the N-dimensional inverse FFT for real input. Args: a (cupy.ndarray): Array to be transform. s (None or tuple of ints): Shape of the output. If ``s`` is not given, they are determined from the lengths of the input along the axes specified by ``axes``. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for transforming ``x`` over ``axes``, which can be obtained using:: plan = cupyx.scipy.fftpack.get_fft_plan(x, s, axes, value_type='C2R') Note that ``plan`` is defaulted to ``None``, meaning CuPy will use an auto-generated plan behind the scene. Returns: cupy.ndarray: The transformed array which shape is specified by ``s`` and type will convert to complex if the input is other. If ``s`` is not given, the length of final transformed axis of output will be ``2*(m-1)`` where `m` is the length of the final transformed axis of the input. .. seealso:: :func:`scipy.fft.irfftn` """ from cupy.cuda import cufft s = _assequence(s) axes = _assequence(axes) func = _default_fft_func(x, s, axes, value_type='C2R') return func(x, s, axes, norm, cufft.CUFFT_INVERSE, 'C2R', overwrite_x=overwrite_x, plan=plan) @_implements(_scipy_fft.hfft) def hfft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the FFT of a signal that has Hermitian symmetry. Args: a (cupy.ndarray): Array to be transform. n (None or int): Length of the transformed axis of the output. For ``n`` output points, ``n//2+1`` input points are necessary. If ``n`` is not given, it is determined from the length of the input along the axis specified by ``axis``. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (None): This argument is currently not supported. Returns: cupy.ndarray: The transformed array which shape is specified by ``n`` and type will convert to complex if the input is other. If ``n`` is not given, the length of the transformed axis is ``2*(m-1)`` where `m` is the length of the transformed axis of the input. .. seealso:: :func:`scipy.fft.hfft` """ # TODO(leofang): support R2C & C2R plans if plan is not None: raise NotImplementedError('hfft plan is currently not yet supported') return _hfft(x, n, axis, norm) @_implements(_scipy_fft.ihfft) def ihfft(x, n=None, axis=-1, norm=None, overwrite_x=False, *, plan=None): """Compute the FFT of a signal that has Hermitian symmetry. Args: a (cupy.ndarray): Array to be transform. n (None or int): Number of points along transformation axis in the input to use. If ``n`` is not given, the length of the input along the axis specified by ``axis`` is used. axis (int): Axis over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. plan (None): This argument is currently not supported. Returns: cupy.ndarray: The transformed array which shape is specified by ``n`` and type will convert to complex if the input is other. The length of the transformed axis is ``n//2+1``. .. seealso:: :func:`scipy.fft.ihfft` """ # TODO(leofang): support R2C & C2R plans if plan is not None: raise NotImplementedError('ihfft plan is currently not yet supported') return _ihfft(x, n, axis, norm) @_implements(_scipy_fft.hfft2) def hfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the FFT of a two-dimensional signal that has Hermitian symmetry. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the real output. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. (This argument is currently not supported) plan (None): This argument is currently not supported. Returns: cupy.ndarray: The real result of the 2-D Hermitian complex real FFT. .. seealso:: :func:`scipy.fft.hfft2` """ if plan is not None: raise NotImplementedError('hfft2 plan is currently not yet supported') return irfft2(x.conj(), s, axes, _swap_direction(norm)) @_implements(_scipy_fft.ihfft2) def ihfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, *, plan=None): """Compute the Inverse FFT of a two-dimensional signal that has Hermitian symmetry. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the real output. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. (This argument is currently not supported) plan (None): This argument is currently not supported. Returns: cupy.ndarray: The real result of the 2-D Hermitian inverse complex real FFT. .. seealso:: :func:`scipy.fft.ihfft2` """ if plan is not None: raise NotImplementedError('ihfft2 plan is currently not yet supported') return rfft2(x, s, axes, _swap_direction(norm)).conj() @_implements(_scipy_fft.hfftn) def hfftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the FFT of a N-dimensional signal that has Hermitian symmetry. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the real output. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. (This argument is currently not supported) plan (None): This argument is currently not supported. Returns: cupy.ndarray: The real result of the N-D Hermitian complex real FFT. .. seealso:: :func:`scipy.fft.hfftn` """ if plan is not None: raise NotImplementedError('hfftn plan is currently not yet supported') return irfftn(x.conj(), s, axes, _swap_direction(norm)) @_implements(_scipy_fft.ihfftn) def ihfftn(x, s=None, axes=None, norm=None, overwrite_x=False, *, plan=None): """Compute the Inverse FFT of a N-dimensional signal that has Hermitian symmetry. Args: x (cupy.ndarray): Array to be transformed. s (None or tuple of ints): Shape of the real output. axes (tuple of ints): Axes over which to compute the FFT. norm (``"backward"``, ``"ortho"``, or ``"forward"``): Optional keyword to specify the normalization mode. Default is ``None``, which is an alias of ``"backward"``. overwrite_x (bool): If True, the contents of ``x`` can be destroyed. (This argument is currently not supported) plan (None): This argument is currently not supported. Returns: cupy.ndarray: The real result of the N-D Hermitian inverse complex real FFT. .. seealso:: :func:`scipy.fft.ihfftn` """ if plan is not None: raise NotImplementedError('ihfftn plan is currently not yet supported') return rfftn(x, s, axes, _swap_direction(norm)).conj()