[Numpy-svn] r8403 - in trunk/numpy/fft: . tests

[numpy-svn at scipy.org]

Author: charris
Date: 2010-05-09 22:41:02 -0500 (Sun, 09 May 2010)
New Revision: 8403

Modified:
   trunk/numpy/fft/helper.py
   trunk/numpy/fft/tests/test_helper.py
Log:
BUG: Make fftshift and ifftshift accept integer arguments for the axes
value. The functions now match their documentation. Fixes ticket #1182,
patch from rgommers.

Modified: trunk/numpy/fft/helper.py
===================================================================
--- trunk/numpy/fft/helper.py	2010-05-09 16:19:08 UTC (rev 8402)
+++ trunk/numpy/fft/helper.py	2010-05-10 03:41:02 UTC (rev 8403)
@@ -7,6 +7,7 @@
 
 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 @@
     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 @@
     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]

Modified: trunk/numpy/fft/tests/test_helper.py
===================================================================
--- trunk/numpy/fft/tests/test_helper.py	2010-05-09 16:19:08 UTC (rev 8402)
+++ trunk/numpy/fft/tests/test_helper.py	2010-05-10 03:41:02 UTC (rev 8403)
@@ -26,6 +26,14 @@
         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):