diff options
Diffstat (limited to 'numpy/fft')
| -rw-r--r-- | numpy/fft/fftpack.py | 24 | ||||
| -rw-r--r-- | numpy/fft/helper.py | 7 |
2 files changed, 31 insertions, 0 deletions
diff --git a/numpy/fft/fftpack.py b/numpy/fft/fftpack.py index e0e96cc79..d88990373 100644 --- a/numpy/fft/fftpack.py +++ b/numpy/fft/fftpack.py @@ -38,6 +38,7 @@ __all__ = ['fft', 'ifft', 'rfft', 'irfft', 'hfft', 'ihfft', 'rfftn', from numpy.core import (array, asarray, zeros, swapaxes, shape, conjugate, take, sqrt) from numpy.core.multiarray import normalize_axis_index +from numpy.core.overrides import array_function_dispatch from . import fftpack_lite as fftpack from .helper import _FFTCache @@ -101,6 +102,11 @@ def _unitary(norm): return norm is not None +def _fft_dispatcher(a, n=None, axis=None, norm=None): + return (a,) + + +@array_function_dispatch(_fft_dispatcher) def fft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional discrete Fourier Transform. @@ -197,6 +203,7 @@ def fft(a, n=None, axis=-1, norm=None): return output +@array_function_dispatch(_fft_dispatcher) def ifft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional inverse discrete Fourier Transform. @@ -290,6 +297,8 @@ def ifft(a, n=None, axis=-1, norm=None): return output * (1 / (sqrt(n) if unitary else n)) + +@array_function_dispatch(_fft_dispatcher) def rfft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional discrete Fourier Transform for real input. @@ -379,6 +388,7 @@ def rfft(a, n=None, axis=-1, norm=None): return output +@array_function_dispatch(_fft_dispatcher) def irfft(a, n=None, axis=-1, norm=None): """ Compute the inverse of the n-point DFT for real input. @@ -469,6 +479,7 @@ def irfft(a, n=None, axis=-1, norm=None): return output * (1 / (sqrt(n) if unitary else n)) +@array_function_dispatch(_fft_dispatcher) def hfft(a, n=None, axis=-1, norm=None): """ Compute the FFT of a signal that has Hermitian symmetry, i.e., a real @@ -551,6 +562,7 @@ def hfft(a, n=None, axis=-1, norm=None): return irfft(conjugate(a), n, axis) * (sqrt(n) if unitary else n) +@array_function_dispatch(_fft_dispatcher) def ihfft(a, n=None, axis=-1, norm=None): """ Compute the inverse FFT of a signal that has Hermitian symmetry. @@ -641,6 +653,11 @@ def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None): return a +def _fftn_dispatcher(a, s=None, axes=None, norm=None): + return (a,) + + +@array_function_dispatch(_fftn_dispatcher) def fftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional discrete Fourier Transform. @@ -738,6 +755,7 @@ def fftn(a, s=None, axes=None, norm=None): return _raw_fftnd(a, s, axes, fft, norm) +@array_function_dispatch(_fftn_dispatcher) def ifftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional inverse discrete Fourier Transform. @@ -835,6 +853,7 @@ def ifftn(a, s=None, axes=None, norm=None): return _raw_fftnd(a, s, axes, ifft, norm) +@array_function_dispatch(_fftn_dispatcher) def fft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional discrete Fourier Transform @@ -925,6 +944,7 @@ def fft2(a, s=None, axes=(-2, -1), norm=None): return _raw_fftnd(a, s, axes, fft, norm) +@array_function_dispatch(_fftn_dispatcher) def ifft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional inverse discrete Fourier Transform. @@ -1012,6 +1032,7 @@ def ifft2(a, s=None, axes=(-2, -1), norm=None): return _raw_fftnd(a, s, axes, ifft, norm) +@array_function_dispatch(_fftn_dispatcher) def rfftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional discrete Fourier Transform for real input. @@ -1104,6 +1125,7 @@ def rfftn(a, s=None, axes=None, norm=None): return a +@array_function_dispatch(_fftn_dispatcher) def rfft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional FFT of a real array. @@ -1141,6 +1163,7 @@ def rfft2(a, s=None, axes=(-2, -1), norm=None): return rfftn(a, s, axes, norm) +@array_function_dispatch(_fftn_dispatcher) def irfftn(a, s=None, axes=None, norm=None): """ Compute the inverse of the N-dimensional FFT of real input. @@ -1235,6 +1258,7 @@ def irfftn(a, s=None, axes=None, norm=None): return a +@array_function_dispatch(_fftn_dispatcher) def irfft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional inverse FFT of a real array. diff --git a/numpy/fft/helper.py b/numpy/fft/helper.py index 729121f31..4b698bb4d 100644 --- a/numpy/fft/helper.py +++ b/numpy/fft/helper.py @@ -11,6 +11,7 @@ except ImportError: import dummy_threading as threading from numpy.compat import integer_types from numpy.core import integer, empty, arange, asarray, roll +from numpy.core.overrides import array_function_dispatch # Created by Pearu Peterson, September 2002 @@ -19,6 +20,11 @@ __all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq'] integer_types = integer_types + (integer,) +def _fftshift_dispatcher(x, axes=None): + return (x,) + + +@array_function_dispatch(_fftshift_dispatcher) def fftshift(x, axes=None): """ Shift the zero-frequency component to the center of the spectrum. @@ -75,6 +81,7 @@ def fftshift(x, axes=None): return roll(x, shift, axes) +@array_function_dispatch(_fftshift_dispatcher) def ifftshift(x, axes=None): """ The inverse of `fftshift`. Although identical for even-length `x`, the |