diff options
Diffstat (limited to 'numpy/fft')
| -rw-r--r-- | numpy/fft/fftpack.py | 2 | ||||
| -rw-r--r-- | numpy/fft/tests/test_helper.py | 45 |
2 files changed, 30 insertions, 17 deletions
diff --git a/numpy/fft/fftpack.py b/numpy/fft/fftpack.py index 2f19ba3d8..80d21ec1d 100644 --- a/numpy/fft/fftpack.py +++ b/numpy/fft/fftpack.py @@ -511,7 +511,7 @@ def _cook_nd_args(a, s=None, axes=None, invreal=0): if len(s) != len(axes): raise ValueError("Shape and axes have different lengths.") if invreal and shapeless: - s[axes[-1]] = (s[axes[-1]] - 1) * 2 + s[-1] = (a.shape[axes[-1]] - 1) * 2 return s, axes diff --git a/numpy/fft/tests/test_helper.py b/numpy/fft/tests/test_helper.py index 8ddac931f..09ccb2b18 100644 --- a/numpy/fft/tests/test_helper.py +++ b/numpy/fft/tests/test_helper.py @@ -3,9 +3,9 @@ """ Test functions for fftpack.helper module """ +import numpy as np from numpy.testing import * -from numpy.fft import fftshift,ifftshift,fftfreq - +from numpy import fft from numpy import pi def random(size): @@ -15,35 +15,48 @@ class TestFFTShift(TestCase): def test_definition(self): x = [0,1,2,3,4,-4,-3,-2,-1] y = [-4,-3,-2,-1,0,1,2,3,4] - assert_array_almost_equal(fftshift(x),y) - assert_array_almost_equal(ifftshift(y),x) + assert_array_almost_equal(fft.fftshift(x),y) + assert_array_almost_equal(fft.ifftshift(y),x) x = [0,1,2,3,4,-5,-4,-3,-2,-1] y = [-5,-4,-3,-2,-1,0,1,2,3,4] - assert_array_almost_equal(fftshift(x),y) - assert_array_almost_equal(ifftshift(y),x) + assert_array_almost_equal(fft.fftshift(x),y) + assert_array_almost_equal(fft.ifftshift(y),x) def test_inverse(self): for n in [1,4,9,100,211]: x = random((n,)) - assert_array_almost_equal(ifftshift(fftshift(x)),x) - + assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)),x) + def test_axes_keyword(self): freqs = [[ 0, 1, 2], [ 3, 4, -4], [-3, -2, -1]] shifted = [[-1, -3, -2], [ 2, 0, 1], [-4, 3, 4]] - assert_array_almost_equal(fftshift(freqs, axes=(0, 1)), shifted) - assert_array_almost_equal(fftshift(freqs, axes=0), fftshift(freqs, axes=(0,))) - assert_array_almost_equal(ifftshift(shifted, axes=(0, 1)), freqs) - assert_array_almost_equal(ifftshift(shifted, axes=0), ifftshift(shifted, axes=(0,))) + assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted) + assert_array_almost_equal(fft.fftshift(freqs, axes=0), + fft.fftshift(freqs, axes=(0,))) + assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs) + assert_array_almost_equal(fft.ifftshift(shifted, axes=0), + fft.ifftshift(shifted, axes=(0,))) class TestFFTFreq(TestCase): def test_definition(self): x = [0,1,2,3,4,-4,-3,-2,-1] - assert_array_almost_equal(9*fftfreq(9),x) - assert_array_almost_equal(9*pi*fftfreq(9,pi),x) + assert_array_almost_equal(9*fft.fftfreq(9),x) + assert_array_almost_equal(9*pi*fft.fftfreq(9,pi),x) x = [0,1,2,3,4,-5,-4,-3,-2,-1] - assert_array_almost_equal(10*fftfreq(10),x) - assert_array_almost_equal(10*pi*fftfreq(10,pi),x) + assert_array_almost_equal(10*fft.fftfreq(10),x) + assert_array_almost_equal(10*pi*fft.fftfreq(10,pi),x) + +class TestIRFFTN(TestCase): + + def test_not_last_axis_success(self): + ar, ai = np.random.random((2, 16, 8, 32)) + a = ar + 1j*ai + + axes = (-2,) + + # Should not raise error + fft.irfftn(a, axes=axes) if __name__ == "__main__": |