Added tests for dst, dstn idst and idstn

wafo/dctpack.py

@@ -351,31 +351,17 @@ def shiftdim(x, n=None):
         return x.reshape((1,) * -n + x.shape)
 
 
-def test_shiftdim():
-    a = np.arange(6).reshape((1, 1, 3, 1, 2))
-
-    print(a.shape)
-    print(a.ndim)
-
-    print(range(a.ndim))
-    # move items 2 places to the left so that x0 <- x2, x1 <- x3, etc
-    print(np.roll(range(a.ndim), -2))
-    print(a.transpose(np.roll(range(a.ndim), -2)))  # transposition of the axes
-    # with a matrix 2x2, A.transpose((1,0)) would be the transpose of A
-    b = shiftdim(a)
-    print(b.shape)
-
-    c = shiftdim(b, -2)
-    print(c.shape)
-
-    print(c == a)
-
-
 def example_dct2():
+    """
+    Example
+    -------
+    >>> example_dct2()
+    """
     import scipy.ndimage as sn
     import matplotlib.pyplot as plt
     name = os.path.join(path, 'autumn.gif')
-    rgb = np.asarray(sn.imread(name), dtype=np.float16)
+    plt.figure(1)
+    rgb = np.asarray(sn.imread(name), dtype=np.float)
     # np.fft.fft(rgb)
     print(np.max(rgb), np.min(rgb))
     plt.set_cmap('jet')
@@ -384,27 +370,25 @@ def example_dct2():
     print(np.abs(rgb-irgb).max())
     plt.imshow(np.log(np.abs(J)))
     # plt.colorbar() #colormap(jet), colorbar
-    plt.show('hold')
+    # plt.show('hold')
+    plt.figure(2)
     v = np.percentile(np.abs(J.ravel()), 10.0)
     print(len(np.flatnonzero(J)), v)
     J[np.abs(J) < v] = 0
     print(len(np.flatnonzero(J)))
     plt.imshow(np.log(np.abs(J)))
-    plt.show('hold')
+    # plt.show('hold')
     K = idctn(J)
     print(np.abs(rgb-K).max())
-    plt.figure(1)
+    plt.figure(3)
     plt.imshow(rgb)
-    plt.figure(2)
+    plt.figure(4)
     plt.imshow(K, vmin=0, vmax=255)
-    plt.show('hold')
-
-
-def test_docstrings():
-    import doctest
-    print('Testing docstrings in %s' % __file__)
-    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
 
 
 if __name__ == '__main__':
-    test_docstrings()
+    from wafo.testing import test_docstrings
+    test_docstrings(__file__)
+    # import matplotlib.pyplot as plt
+    # example_dct2()
+    # plt.show('hold')

wafo/tests/test_dct_pack.py

@@ -18,7 +18,7 @@ class Test(unittest.TestCase):
         assert_array_almost_equal(c.shape, a.shape)
         assert_array_almost_equal(c, a)
 
-    def test_dct3(self):
+    def _test_3d(self, dct, idct, dctn, idctn):
         a = np.array([[[0.51699637,  0.42946223,  0.89843545],
                        [0.27853391,  0.8931508,  0.34319118],
                        [0.51984431,  0.09217771,  0.78764716]],
@@ -28,26 +28,37 @@ class Test(unittest.TestCase):
                       [[0.86098829,  0.07901332,  0.82169182],
                        [0.12560088,  0.78210188,  0.69805434],
                        [0.33544628,  0.81540172,  0.9393219]]])
-        dct = wd.dct
         d = dct(dct(dct(a).transpose(0, 2, 1)).transpose(2, 1, 0)
                 ).transpose(2, 1, 0).transpose(0, 2, 1)
-        d0 = wd.dctn(a)
+        d0 = dctn(a)
-        idct = wd.idct
+        e0 = idctn(d0)
         e = idct(idct(idct(d).transpose(0, 2, 1)).transpose(2, 1, 0)
                  ).transpose(2, 1, 0).transpose(0, 2, 1)
         assert_array_almost_equal(d, d0)
+        assert_array_almost_equal(e, e0)
         assert_array_almost_equal(a, e)
 
+    def test_3d_dct_idct_dctn_idctn(self):
+        self._test_3d(wd.dct, wd.idct, wd.dctn, wd.idctn)
+
+    def test_3d_dst_idst_dstn_idstn(self):
+        self._test_3d(wd.dst, wd.idst, wd.dstn, wd.idstn)
+
     def test_dct_and_dctn(self):
-        a = np.arange(12).reshape((3, -1))
+        self._test_1d(wd.dct, wd.idct, wd.dctn, wd.idctn)
 
-        y = wd.dct(a)
+    def test_dst_and_dstn(self):
-        x = wd.idct(y)
+        self._test_1d(wd.dst, wd.idst, wd.dstn, wd.idstn)
-        assert_array_almost_equal(x, a)
 
-        yn = wd.dctn(a)  # , shape=(10,), axes=(1,))
+    def _test_1d(self, dct, idct, dctn, idctn):
-        xn = wd.idctn(yn)  # , axes=(1,))
+        a = np.arange(12)  # .reshape((3, -1))
+        y = dct(a)
+        yn = dctn(a)
+        x = idct(y)
+        xn = idctn(yn)
 
+        assert_array_almost_equal(y, yn)
+        assert_array_almost_equal(x, a)
         assert_array_almost_equal(xn, a)