[Scipy-svn] r5470 - in trunk/scipy/fftpack: . tests

[scipy-svn at scipy.org]

Author: cdavid
Date: 2009-01-16 01:34:55 -0600 (Fri, 16 Jan 2009)
New Revision: 5470

Modified:
   trunk/scipy/fftpack/realtransforms.py
   trunk/scipy/fftpack/tests/test_real_transforms.py
Log:
Add support for the axis argument in dct.

Modified: trunk/scipy/fftpack/realtransforms.py
===================================================================
--- trunk/scipy/fftpack/realtransforms.py	2009-01-16 07:34:39 UTC (rev 5469)
+++ trunk/scipy/fftpack/realtransforms.py	2009-01-16 07:34:55 UTC (rev 5470)
@@ -11,7 +11,7 @@
 atexit.register(_fftpack.destroy_dct1_cache)
 atexit.register(_fftpack.destroy_dct2_cache)
 
-def dct1(x, n=None):
+def dct1(x, n=None, axis=-1):
     """
     Return Discrete Cosine Transform (type I) of arbitrary type sequence x.
 
@@ -21,14 +21,16 @@
         input array.
     n : int, optional
         Length of the transform.
+    axis : int, optional
+        axis over which to compute the transform.
 
     Returns
     -------
     y : real ndarray
     """
-    return _dct(x, 1, n)
+    return _dct(x, 1, n, axis)
 
-def dct2(x, n=None):
+def dct2(x, n=None, axis=-1):
     """
     Return Discrete Cosine Transform (type II) of arbitrary type sequence x.
     There are several definitions, we use the following:
@@ -45,6 +47,8 @@
         input array.
     n : int, optional
         Length of the transform.
+    axis : int, optional
+        axis over which to compute the transform.
 
     Returns
     -------
@@ -58,7 +62,7 @@
     'A Fast Cosine Transform in One and Two Dimensions', by J. Makhoul, in IEEE
     Transactions on acoustics, speech and signal processing.
     """
-    return _dct(x, 2, n)
+    return _dct(x, 2, n, axis)
 
 def _dct(x, type, n=None, axis=-1, overwrite_x=0):
     """
@@ -98,9 +102,9 @@
 
     if axis == -1 or axis == len(tmp.shape) - 1:
         return f(tmp, n, 0, overwrite_x)
-    else:
-        raise NotImplementedError("Axis arg not yet implemented")
+    #else:
+    #    raise NotImplementedError("Axis arg not yet implemented")
 
-    #tmp = swapaxes(tmp, axis, -1)
-    #tmp = work_function(tmp,n,1,0,overwrite_x)
-    #return swapaxes(tmp, axis, -1)
+    tmp = np.swapaxes(tmp, axis, -1)
+    tmp = f(tmp, n, 0, overwrite_x)
+    return np.swapaxes(tmp, axis, -1)

Modified: trunk/scipy/fftpack/tests/test_real_transforms.py
===================================================================
--- trunk/scipy/fftpack/tests/test_real_transforms.py	2009-01-16 07:34:39 UTC (rev 5469)
+++ trunk/scipy/fftpack/tests/test_real_transforms.py	2009-01-16 07:34:55 UTC (rev 5470)
@@ -122,6 +122,19 @@
                     "Output dtype is %s, expected %s" % (y.dtype, self.rdt))
             assert_array_almost_equal(y, yr)
 
+    def test_axis(self):
+        nt = 2
+        for i in [7, 8, 9, 16, 32, 64]:
+            x = np.random.randn(nt, i)
+            y = dct2(x)
+            for j in range(nt):
+                assert_array_almost_equal(y[j], dct2(x[j]))
+
+            x = x.T
+            y = dct2(x, axis=0)
+            for j in range(nt):
+                assert_array_almost_equal(y[:,j], dct2(x[:,j]))
+
 class TestDCTIIDouble(_TestDCTIIBase):
     def setUp(self):
         self.rdt = np.double