Chris@87: #!/usr/bin/env python
Chris@87: """Test functions for fftpack.helper module
Chris@87: 
Chris@87: Copied from fftpack.helper by Pearu Peterson, October 2005
Chris@87: 
Chris@87: """
Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: import numpy as np
Chris@87: from numpy.testing import TestCase, run_module_suite, assert_array_almost_equal
Chris@87: from numpy import fft
Chris@87: from numpy import pi
Chris@87: 
Chris@87: 
Chris@87: class TestFFTShift(TestCase):
Chris@87: 
Chris@87:     def test_definition(self):
Chris@87:         x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
Chris@87:         y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
Chris@87:         assert_array_almost_equal(fft.fftshift(x), y)
Chris@87:         assert_array_almost_equal(fft.ifftshift(y), x)
Chris@87:         x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
Chris@87:         y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
Chris@87:         assert_array_almost_equal(fft.fftshift(x), y)
Chris@87:         assert_array_almost_equal(fft.ifftshift(y), x)
Chris@87: 
Chris@87:     def test_inverse(self):
Chris@87:         for n in [1, 4, 9, 100, 211]:
Chris@87:             x = np.random.random((n,))
Chris@87:             assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)), x)
Chris@87: 
Chris@87:     def test_axes_keyword(self):
Chris@87:         freqs = [[ 0,  1,  2], [ 3,  4, -4], [-3, -2, -1]]
Chris@87:         shifted = [[-1, -3, -2], [ 2,  0,  1], [-4,  3,  4]]
Chris@87:         assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted)
Chris@87:         assert_array_almost_equal(fft.fftshift(freqs, axes=0),
Chris@87:                 fft.fftshift(freqs, axes=(0,)))
Chris@87:         assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs)
Chris@87:         assert_array_almost_equal(fft.ifftshift(shifted, axes=0),
Chris@87:                 fft.ifftshift(shifted, axes=(0,)))
Chris@87: 
Chris@87: 
Chris@87: class TestFFTFreq(TestCase):
Chris@87: 
Chris@87:     def test_definition(self):
Chris@87:         x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
Chris@87:         assert_array_almost_equal(9*fft.fftfreq(9), x)
Chris@87:         assert_array_almost_equal(9*pi*fft.fftfreq(9, pi), x)
Chris@87:         x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
Chris@87:         assert_array_almost_equal(10*fft.fftfreq(10), x)
Chris@87:         assert_array_almost_equal(10*pi*fft.fftfreq(10, pi), x)
Chris@87: 
Chris@87: 
Chris@87: class TestRFFTFreq(TestCase):
Chris@87: 
Chris@87:     def test_definition(self):
Chris@87:         x = [0, 1, 2, 3, 4]
Chris@87:         assert_array_almost_equal(9*fft.rfftfreq(9), x)
Chris@87:         assert_array_almost_equal(9*pi*fft.rfftfreq(9, pi), x)
Chris@87:         x = [0, 1, 2, 3, 4, 5]
Chris@87:         assert_array_almost_equal(10*fft.rfftfreq(10), x)
Chris@87:         assert_array_almost_equal(10*pi*fft.rfftfreq(10, pi), x)
Chris@87: 
Chris@87: 
Chris@87: class TestIRFFTN(TestCase):
Chris@87: 
Chris@87:     def test_not_last_axis_success(self):
Chris@87:         ar, ai = np.random.random((2, 16, 8, 32))
Chris@87:         a = ar + 1j*ai
Chris@87: 
Chris@87:         axes = (-2,)
Chris@87: 
Chris@87:         # Should not raise error
Chris@87:         fft.irfftn(a, axes=axes)
Chris@87: 
Chris@87: 
Chris@87: if __name__ == "__main__":
Chris@87:     run_module_suite()