Merge pull request #5992 from charris/add-norm-keyword-to-fft-functions

ENH: Add a norm keyword and tests for fft transforms

4 files changed, 222 insertions, 43 deletions

diff --git a/doc/release/1.10.0-notes.rst b/doc/release/1.10.0-notes.rst
index b24082975..c75139dc7 100644
--- a/doc/release/1.10.0-notes.rst
+++ b/doc/release/1.10.0-notes.rst
@@ -198,6 +198,14 @@ equivalent function ``matmul`` has also been added for testing purposes and
 use in earlier Python versions. The function is preliminary and the order
 and number of its optional arguments can be expected to change.
 
+New argument ``norm`` to fft functions
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+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}`.
+
 
 Improvements
 ============
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):