diff options
Diffstat (limited to 'numpy/fft')
-rw-r--r-- | numpy/fft/fftpack.py | 150 | ||||
-rw-r--r-- | numpy/fft/info.py | 10 | ||||
-rw-r--r-- | numpy/fft/tests/test_fftpack.py | 97 |
3 files changed, 214 insertions, 43 deletions
diff --git a/numpy/fft/fftpack.py b/numpy/fft/fftpack.py index e12ae1eec..4ad4f6802 100644 --- a/numpy/fft/fftpack.py +++ b/numpy/fft/fftpack.py @@ -35,7 +35,8 @@ from __future__ import division, absolute_import, print_function __all__ = ['fft', 'ifft', 'rfft', 'irfft', 'hfft', 'ihfft', 'rfftn', 'irfftn', 'rfft2', 'irfft2', 'fft2', 'ifft2', 'fftn', 'ifftn'] -from numpy.core import array, asarray, zeros, swapaxes, shape, conjugate, take +from numpy.core import (array, asarray, zeros, swapaxes, shape, conjugate, + take, sqrt) from . import fftpack_lite as fftpack _fft_cache = {} @@ -50,7 +51,8 @@ def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, n = a.shape[axis] if n < 1: - raise ValueError("Invalid number of FFT data points (%d) specified." % n) + raise ValueError("Invalid number of FFT data points (%d) specified." + % n) try: # Thread-safety note: We rely on list.pop() here to atomically @@ -88,7 +90,14 @@ def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, return r -def fft(a, n=None, axis=-1): +def _unitary(norm): + if norm not in (None, "ortho"): + raise ValueError("Invalid norm value %s, should be None or \"ortho\"." + % norm) + return norm is not None + + +def fft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional discrete Fourier Transform. @@ -108,6 +117,9 @@ def fft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the FFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -171,10 +183,16 @@ def fft(a, n=None, axis=-1): """ - return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftf, _fft_cache) + a = asarray(a).astype(complex) + if n is None: + n = a.shape[axis] + output = _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftf, _fft_cache) + if _unitary(norm): + output *= 1 / sqrt(n) + return output -def ifft(a, n=None, axis=-1): +def ifft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional inverse discrete Fourier Transform. @@ -203,6 +221,9 @@ def ifft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the inverse DFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -251,11 +272,13 @@ def ifft(a, n=None, axis=-1): # The copy may be required for multithreading. a = array(a, copy=True, dtype=complex) if n is None: - n = shape(a)[axis] - return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftb, _fft_cache) / n + n = a.shape[axis] + unitary = _unitary(norm) + output = _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftb, _fft_cache) + return output * (1 / (sqrt(n) if unitary else n)) -def rfft(a, n=None, axis=-1): +def rfft(a, n=None, axis=-1, norm=None): """ Compute the one-dimensional discrete Fourier Transform for real input. @@ -275,6 +298,9 @@ def rfft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the FFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -331,10 +357,14 @@ def rfft(a, n=None, axis=-1): """ # The copy may be required for multithreading. a = array(a, copy=True, dtype=float) - return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftf, _real_fft_cache) + output = _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftf, + _real_fft_cache) + if _unitary(norm): + output *= 1 / sqrt(a.shape[axis]) + return output -def irfft(a, n=None, axis=-1): +def irfft(a, n=None, axis=-1, norm=None): """ Compute the inverse of the n-point DFT for real input. @@ -362,6 +392,9 @@ def irfft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the inverse FFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -413,12 +446,14 @@ def irfft(a, n=None, axis=-1): # The copy may be required for multithreading. a = array(a, copy=True, dtype=complex) if n is None: - n = (shape(a)[axis] - 1) * 2 - return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftb, - _real_fft_cache) / n + n = (a.shape[axis] - 1) * 2 + unitary = _unitary(norm) + output = _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftb, + _real_fft_cache) + return output * (1 / (sqrt(n) if unitary else n)) -def hfft(a, n=None, axis=-1): +def hfft(a, n=None, axis=-1, norm=None): """ Compute the FFT of a signal which has Hermitian symmetry (real spectrum). @@ -435,6 +470,9 @@ def hfft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the FFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -487,11 +525,12 @@ def hfft(a, n=None, axis=-1): # The copy may be required for multithreading. a = array(a, copy=True, dtype=complex) if n is None: - n = (shape(a)[axis] - 1) * 2 - return irfft(conjugate(a), n, axis) * n + n = (a.shape[axis] - 1) * 2 + unitary = _unitary(norm) + return irfft(conjugate(a), n, axis) * (sqrt(n) if unitary else n) -def ihfft(a, n=None, axis=-1): +def ihfft(a, n=None, axis=-1, norm=None): """ Compute the inverse FFT of a signal which has Hermitian symmetry. @@ -508,6 +547,9 @@ def ihfft(a, n=None, axis=-1): axis : int, optional Axis over which to compute the inverse FFT. If not given, the last axis is used. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -541,8 +583,10 @@ def ihfft(a, n=None, axis=-1): # The copy may be required for multithreading. a = array(a, copy=True, dtype=float) if n is None: - n = shape(a)[axis] - return conjugate(rfft(a, n, axis))/n + n = a.shape[axis] + unitary = _unitary(norm) + output = conjugate(rfft(a, n, axis)) + return output * (1 / (sqrt(n) if unitary else n)) def _cook_nd_args(a, s=None, axes=None, invreal=0): @@ -564,17 +608,17 @@ def _cook_nd_args(a, s=None, axes=None, invreal=0): return s, axes -def _raw_fftnd(a, s=None, axes=None, function=fft): +def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None): a = asarray(a) s, axes = _cook_nd_args(a, s, axes) itl = list(range(len(axes))) itl.reverse() for ii in itl: - a = function(a, n=s[ii], axis=axes[ii]) + a = function(a, n=s[ii], axis=axes[ii], norm=norm) return a -def fftn(a, s=None, axes=None): +def fftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional discrete Fourier Transform. @@ -599,6 +643,9 @@ def fftn(a, s=None, axes=None): axes are used, or all axes if `s` is also not specified. Repeated indices in `axes` means that the transform over that axis is performed multiple times. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -664,9 +711,10 @@ def fftn(a, s=None, axes=None): """ - return _raw_fftnd(a, s, axes, fft) + return _raw_fftnd(a, s, axes, fft, norm) + -def ifftn(a, s=None, axes=None): +def ifftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional inverse discrete Fourier Transform. @@ -700,6 +748,9 @@ def ifftn(a, s=None, axes=None): axes are used, or all axes if `s` is also not specified. Repeated indices in `axes` means that the inverse transform over that axis is performed multiple times. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -756,10 +807,10 @@ def ifftn(a, s=None, axes=None): """ - return _raw_fftnd(a, s, axes, ifft) + return _raw_fftnd(a, s, axes, ifft, norm) -def fft2(a, s=None, axes=(-2, -1)): +def fft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional discrete Fourier Transform @@ -785,6 +836,9 @@ def fft2(a, s=None, axes=(-2, -1)): axes are used. A repeated index in `axes` means the transform over that axis is performed multiple times. A one-element sequence means that a one-dimensional FFT is performed. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -842,10 +896,10 @@ def fft2(a, s=None, axes=(-2, -1)): """ - return _raw_fftnd(a, s, axes, fft) + return _raw_fftnd(a, s, axes, fft, norm) -def ifft2(a, s=None, axes=(-2, -1)): +def ifft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional inverse discrete Fourier Transform. @@ -878,6 +932,9 @@ def ifft2(a, s=None, axes=(-2, -1)): axes are used. A repeated index in `axes` means the transform over that axis is performed multiple times. A one-element sequence means that a one-dimensional FFT is performed. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -925,10 +982,10 @@ def ifft2(a, s=None, axes=(-2, -1)): """ - return _raw_fftnd(a, s, axes, ifft) + return _raw_fftnd(a, s, axes, ifft, norm) -def rfftn(a, s=None, axes=None): +def rfftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional discrete Fourier Transform for real input. @@ -954,6 +1011,9 @@ def rfftn(a, s=None, axes=None): axes : sequence of ints, optional Axes over which to compute the FFT. If not given, the last ``len(s)`` axes are used, or all axes if `s` is also not specified. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -1010,12 +1070,13 @@ def rfftn(a, s=None, axes=None): # The copy may be required for multithreading. a = array(a, copy=True, dtype=float) s, axes = _cook_nd_args(a, s, axes) - a = rfft(a, s[-1], axes[-1]) + a = rfft(a, s[-1], axes[-1], norm) for ii in range(len(axes)-1): - a = fft(a, s[ii], axes[ii]) + a = fft(a, s[ii], axes[ii], norm) return a -def rfft2(a, s=None, axes=(-2, -1)): + +def rfft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional FFT of a real array. @@ -1027,6 +1088,9 @@ def rfft2(a, s=None, axes=(-2, -1)): Shape of the FFT. axes : sequence of ints, optional Axes over which to compute the FFT. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -1045,9 +1109,10 @@ def rfft2(a, s=None, axes=(-2, -1)): """ - return rfftn(a, s, axes) + return rfftn(a, s, axes, norm) -def irfftn(a, s=None, axes=None): + +def irfftn(a, s=None, axes=None, norm=None): """ Compute the inverse of the N-dimensional FFT of real input. @@ -1080,6 +1145,9 @@ def irfftn(a, s=None, axes=None): `len(s)` axes are used, or all axes if `s` is also not specified. Repeated indices in `axes` means that the inverse transform over that axis is performed multiple times. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -1132,11 +1200,12 @@ def irfftn(a, s=None, axes=None): a = array(a, copy=True, dtype=complex) s, axes = _cook_nd_args(a, s, axes, invreal=1) for ii in range(len(axes)-1): - a = ifft(a, s[ii], axes[ii]) - a = irfft(a, s[-1], axes[-1]) + a = ifft(a, s[ii], axes[ii], norm) + a = irfft(a, s[-1], axes[-1], norm) return a -def irfft2(a, s=None, axes=(-2, -1)): + +def irfft2(a, s=None, axes=(-2, -1), norm=None): """ Compute the 2-dimensional inverse FFT of a real array. @@ -1149,6 +1218,9 @@ def irfft2(a, s=None, axes=(-2, -1)): axes : sequence of ints, optional The axes over which to compute the inverse fft. Default is the last two axes. + norm : {None, "ortho"}, optional + .. versionadded:: 1.10.0 + Normalization mode (see `numpy.fft`). Default is None. Returns ------- @@ -1166,4 +1238,4 @@ def irfft2(a, s=None, axes=(-2, -1)): """ - return irfftn(a, s, axes) + return irfftn(a, s, axes, norm) diff --git a/numpy/fft/info.py b/numpy/fft/info.py index 916d452f2..f7fc95453 100644 --- a/numpy/fft/info.py +++ b/numpy/fft/info.py @@ -116,7 +116,15 @@ The inverse DFT is defined as \\qquad m = 0,\\ldots,n-1. It differs from the forward transform by the sign of the exponential -argument and the normalization by :math:`1/n`. +argument and the default normalization by :math:`1/n`. + +Normalization +------------- +The default normalization has the direct transforms unscaled and the inverse +transforms are scaled by :math:`1/n`. It is possible to obtain unitary +transforms by setting the keyword argument ``norm`` to ``"ortho"`` (default is +`None`) so that both direct and inverse transforms will be scaled by +:math:`1/\\sqrt{n}`. Real and Hermitian transforms ----------------------------- diff --git a/numpy/fft/tests/test_fftpack.py b/numpy/fft/tests/test_fftpack.py index 45b5ac784..2e6294252 100644 --- a/numpy/fft/tests/test_fftpack.py +++ b/numpy/fft/tests/test_fftpack.py @@ -1,6 +1,7 @@ from __future__ import division, absolute_import, print_function import numpy as np +from numpy.random import random from numpy.testing import TestCase, run_module_suite, assert_array_almost_equal from numpy.testing import assert_array_equal import threading @@ -26,10 +27,100 @@ class TestFFTShift(TestCase): class TestFFT1D(TestCase): - def test_basic(self): - rand = np.random.random - x = rand(30) + 1j*rand(30) + def test_fft(self): + x = random(30) + 1j*random(30) assert_array_almost_equal(fft1(x), np.fft.fft(x)) + assert_array_almost_equal(fft1(x) / np.sqrt(30), + np.fft.fft(x, norm="ortho")) + + def test_ifft(self): + x = random(30) + 1j*random(30) + assert_array_almost_equal(x, np.fft.ifft(np.fft.fft(x))) + assert_array_almost_equal( + x, np.fft.ifft(np.fft.fft(x, norm="ortho"), norm="ortho")) + + def test_fft2(self): + x = random((30, 20)) + 1j*random((30, 20)) + assert_array_almost_equal(np.fft.fft(np.fft.fft(x, axis=1), axis=0), + np.fft.fft2(x)) + assert_array_almost_equal(np.fft.fft2(x) / np.sqrt(30 * 20), + np.fft.fft2(x, norm="ortho")) + + def test_ifft2(self): + x = random((30, 20)) + 1j*random((30, 20)) + assert_array_almost_equal(np.fft.ifft(np.fft.ifft(x, axis=1), axis=0), + np.fft.ifft2(x)) + assert_array_almost_equal(np.fft.ifft2(x) * np.sqrt(30 * 20), + np.fft.ifft2(x, norm="ortho")) + + def test_fftn(self): + x = random((30, 20, 10)) + 1j*random((30, 20, 10)) + assert_array_almost_equal( + np.fft.fft(np.fft.fft(np.fft.fft(x, axis=2), axis=1), axis=0), + np.fft.fftn(x)) + assert_array_almost_equal(np.fft.fftn(x) / np.sqrt(30 * 20 * 10), + np.fft.fftn(x, norm="ortho")) + + def test_ifftn(self): + x = random((30, 20, 10)) + 1j*random((30, 20, 10)) + assert_array_almost_equal( + np.fft.ifft(np.fft.ifft(np.fft.ifft(x, axis=2), axis=1), axis=0), + np.fft.ifftn(x)) + assert_array_almost_equal(np.fft.ifftn(x) * np.sqrt(30 * 20 * 10), + np.fft.ifftn(x, norm="ortho")) + + def test_rfft(self): + x = random(30) + assert_array_almost_equal(np.fft.fft(x)[:16], np.fft.rfft(x)) + assert_array_almost_equal(np.fft.rfft(x) / np.sqrt(30), + np.fft.rfft(x, norm="ortho")) + + def test_irfft(self): + x = random(30) + assert_array_almost_equal(x, np.fft.irfft(np.fft.rfft(x))) + assert_array_almost_equal( + x, np.fft.irfft(np.fft.rfft(x, norm="ortho"), norm="ortho")) + + def test_rfft2(self): + x = random((30, 20)) + assert_array_almost_equal(np.fft.fft2(x)[:, :11], np.fft.rfft2(x)) + assert_array_almost_equal(np.fft.rfft2(x) / np.sqrt(30 * 20), + np.fft.rfft2(x, norm="ortho")) + + def test_irfft2(self): + x = random((30, 20)) + assert_array_almost_equal(x, np.fft.irfft2(np.fft.rfft2(x))) + assert_array_almost_equal( + x, np.fft.irfft2(np.fft.rfft2(x, norm="ortho"), norm="ortho")) + + def test_rfftn(self): + x = random((30, 20, 10)) + assert_array_almost_equal(np.fft.fftn(x)[:, :, :6], np.fft.rfftn(x)) + assert_array_almost_equal(np.fft.rfftn(x) / np.sqrt(30 * 20 * 10), + np.fft.rfftn(x, norm="ortho")) + + def test_irfftn(self): + x = random((30, 20, 10)) + assert_array_almost_equal(x, np.fft.irfftn(np.fft.rfftn(x))) + assert_array_almost_equal( + x, np.fft.irfftn(np.fft.rfftn(x, norm="ortho"), norm="ortho")) + + def test_hfft(self): + x = random(14) + 1j*random(14) + x_herm = np.concatenate((random(1), x, random(1))) + x = np.concatenate((x_herm, x[::-1].conj())) + assert_array_almost_equal(np.fft.fft(x), np.fft.hfft(x_herm)) + assert_array_almost_equal(np.fft.hfft(x_herm) / np.sqrt(30), + np.fft.hfft(x_herm, norm="ortho")) + + def test_ihttf(self): + x = random(14) + 1j*random(14) + x_herm = np.concatenate((random(1), x, random(1))) + x = np.concatenate((x_herm, x[::-1].conj())) + assert_array_almost_equal(x_herm, np.fft.ihfft(np.fft.hfft(x_herm))) + assert_array_almost_equal( + x_herm, np.fft.ihfft(np.fft.hfft(x_herm, norm="ortho"), + norm="ortho")) class TestFFTThreadSafe(TestCase): |