diff options
Diffstat (limited to 'numpy/fft/fftpack.py')
| -rw-r--r-- | numpy/fft/fftpack.py | 30 |
1 files changed, 28 insertions, 2 deletions
diff --git a/numpy/fft/fftpack.py b/numpy/fft/fftpack.py index e17e1cb34..d88990373 100644 --- a/numpy/fft/fftpack.py +++ b/numpy/fft/fftpack.py @@ -37,6 +37,8 @@ __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 @@ -47,6 +49,7 @@ _real_fft_cache = _FFTCache(max_size_in_mb=100, max_item_count=32) def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, work_function=fftpack.cfftf, fft_cache=_fft_cache): a = asarray(a) + axis = normalize_axis_index(axis, a.ndim) if n is None: n = a.shape[axis] @@ -78,10 +81,10 @@ def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, z[tuple(index)] = a a = z - if axis != -1: + if axis != a.ndim - 1: a = swapaxes(a, axis, -1) r = work_function(a, wsave) - if axis != -1: + if axis != a.ndim - 1: r = swapaxes(r, axis, -1) # As soon as we put wsave back into the cache, another thread could pick it @@ -99,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. @@ -195,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. @@ -288,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. @@ -377,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. @@ -467,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 @@ -549,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. @@ -639,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. @@ -736,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. @@ -833,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 @@ -923,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. @@ -1010,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. @@ -1102,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. @@ -1139,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. @@ -1233,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. |