Chris@87
|
1 #!/usr/bin/env python
|
Chris@87
|
2 """Test functions for fftpack.helper module
|
Chris@87
|
3
|
Chris@87
|
4 Copied from fftpack.helper by Pearu Peterson, October 2005
|
Chris@87
|
5
|
Chris@87
|
6 """
|
Chris@87
|
7 from __future__ import division, absolute_import, print_function
|
Chris@87
|
8
|
Chris@87
|
9 import numpy as np
|
Chris@87
|
10 from numpy.testing import TestCase, run_module_suite, assert_array_almost_equal
|
Chris@87
|
11 from numpy import fft
|
Chris@87
|
12 from numpy import pi
|
Chris@87
|
13
|
Chris@87
|
14
|
Chris@87
|
15 class TestFFTShift(TestCase):
|
Chris@87
|
16
|
Chris@87
|
17 def test_definition(self):
|
Chris@87
|
18 x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
|
Chris@87
|
19 y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
|
Chris@87
|
20 assert_array_almost_equal(fft.fftshift(x), y)
|
Chris@87
|
21 assert_array_almost_equal(fft.ifftshift(y), x)
|
Chris@87
|
22 x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
|
Chris@87
|
23 y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
|
Chris@87
|
24 assert_array_almost_equal(fft.fftshift(x), y)
|
Chris@87
|
25 assert_array_almost_equal(fft.ifftshift(y), x)
|
Chris@87
|
26
|
Chris@87
|
27 def test_inverse(self):
|
Chris@87
|
28 for n in [1, 4, 9, 100, 211]:
|
Chris@87
|
29 x = np.random.random((n,))
|
Chris@87
|
30 assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)), x)
|
Chris@87
|
31
|
Chris@87
|
32 def test_axes_keyword(self):
|
Chris@87
|
33 freqs = [[ 0, 1, 2], [ 3, 4, -4], [-3, -2, -1]]
|
Chris@87
|
34 shifted = [[-1, -3, -2], [ 2, 0, 1], [-4, 3, 4]]
|
Chris@87
|
35 assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted)
|
Chris@87
|
36 assert_array_almost_equal(fft.fftshift(freqs, axes=0),
|
Chris@87
|
37 fft.fftshift(freqs, axes=(0,)))
|
Chris@87
|
38 assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs)
|
Chris@87
|
39 assert_array_almost_equal(fft.ifftshift(shifted, axes=0),
|
Chris@87
|
40 fft.ifftshift(shifted, axes=(0,)))
|
Chris@87
|
41
|
Chris@87
|
42
|
Chris@87
|
43 class TestFFTFreq(TestCase):
|
Chris@87
|
44
|
Chris@87
|
45 def test_definition(self):
|
Chris@87
|
46 x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
|
Chris@87
|
47 assert_array_almost_equal(9*fft.fftfreq(9), x)
|
Chris@87
|
48 assert_array_almost_equal(9*pi*fft.fftfreq(9, pi), x)
|
Chris@87
|
49 x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
|
Chris@87
|
50 assert_array_almost_equal(10*fft.fftfreq(10), x)
|
Chris@87
|
51 assert_array_almost_equal(10*pi*fft.fftfreq(10, pi), x)
|
Chris@87
|
52
|
Chris@87
|
53
|
Chris@87
|
54 class TestRFFTFreq(TestCase):
|
Chris@87
|
55
|
Chris@87
|
56 def test_definition(self):
|
Chris@87
|
57 x = [0, 1, 2, 3, 4]
|
Chris@87
|
58 assert_array_almost_equal(9*fft.rfftfreq(9), x)
|
Chris@87
|
59 assert_array_almost_equal(9*pi*fft.rfftfreq(9, pi), x)
|
Chris@87
|
60 x = [0, 1, 2, 3, 4, 5]
|
Chris@87
|
61 assert_array_almost_equal(10*fft.rfftfreq(10), x)
|
Chris@87
|
62 assert_array_almost_equal(10*pi*fft.rfftfreq(10, pi), x)
|
Chris@87
|
63
|
Chris@87
|
64
|
Chris@87
|
65 class TestIRFFTN(TestCase):
|
Chris@87
|
66
|
Chris@87
|
67 def test_not_last_axis_success(self):
|
Chris@87
|
68 ar, ai = np.random.random((2, 16, 8, 32))
|
Chris@87
|
69 a = ar + 1j*ai
|
Chris@87
|
70
|
Chris@87
|
71 axes = (-2,)
|
Chris@87
|
72
|
Chris@87
|
73 # Should not raise error
|
Chris@87
|
74 fft.irfftn(a, axes=axes)
|
Chris@87
|
75
|
Chris@87
|
76
|
Chris@87
|
77 if __name__ == "__main__":
|
Chris@87
|
78 run_module_suite()
|