diff options
| author | Charles Harris <charlesr.harris@gmail.com> | 2010-05-10 03:41:02 +0000 |
|---|---|---|
| committer | Charles Harris <charlesr.harris@gmail.com> | 2010-05-10 03:41:02 +0000 |
| commit | 427a9632db6ddc1e8659ab3191fe63873a8a188b (patch) | |
| tree | c79dd77a275676ad62c9f3fd0d043044884a91a5 /numpy/fft | |
| parent | 81a86c775143d7335613fdb960d6267fc3d602f8 (diff) | |
| download | numpy-427a9632db6ddc1e8659ab3191fe63873a8a188b.tar.gz | |
BUG: Make fftshift and ifftshift accept integer arguments for the axes
value. The functions now match their documentation. Fixes ticket #1182,
patch from rgommers.
Diffstat (limited to 'numpy/fft')
| -rw-r--r-- | numpy/fft/helper.py | 5 | ||||
| -rw-r--r-- | numpy/fft/tests/test_helper.py | 8 |
2 files changed, 13 insertions, 0 deletions
diff --git a/numpy/fft/helper.py b/numpy/fft/helper.py index 8bd2564de..f6c570445 100644 --- a/numpy/fft/helper.py +++ b/numpy/fft/helper.py @@ -7,6 +7,7 @@ __all__ = ['fftshift','ifftshift','fftfreq'] from numpy.core import asarray, concatenate, arange, take, \ integer, empty +import numpy.core.numerictypes as nt import types def fftshift(x,axes=None): @@ -57,6 +58,8 @@ def fftshift(x,axes=None): ndim = len(tmp.shape) if axes is None: axes = range(ndim) + elif isinstance(axes, (int, nt.integer)): + axes = (axes,) y = tmp for k in axes: n = tmp.shape[k] @@ -103,6 +106,8 @@ def ifftshift(x,axes=None): ndim = len(tmp.shape) if axes is None: axes = range(ndim) + elif isinstance(axes, (int, nt.integer)): + axes = (axes,) y = tmp for k in axes: n = tmp.shape[k] diff --git a/numpy/fft/tests/test_helper.py b/numpy/fft/tests/test_helper.py index f757b6032..8ddac931f 100644 --- a/numpy/fft/tests/test_helper.py +++ b/numpy/fft/tests/test_helper.py @@ -26,6 +26,14 @@ class TestFFTShift(TestCase): for n in [1,4,9,100,211]: x = random((n,)) assert_array_almost_equal(ifftshift(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,))) class TestFFTFreq(TestCase): |