TST: add tests for pywt.pad

Author: grlee77

pywt/_dwt.py

@@ -434,11 +434,14 @@ def pad(x, pad_widths, mode):
     for modes `smooth` and `antisymmetric` as these modes are not supported in
     an efficient manner by the underlying `numpy.pad` function.
     """
-    if np.isscalar(pad_widths):
-        pad_widths = (pad_widths, )
-    if len(pad_widths) == 1:
-        pad_widths = (pad_widths[0], ) * x.ndim
-    pad_widths = [(p, p) if np.isscalar(p) else p for p in pad_widths]
+    x = np.asanyarray(x)
+
+    # process pad_widths exactly as in numpy.pad
+    pad_widths = np.array(pad_widths)
+    pad_widths = np.round(pad_widths).astype(np.intp, copy=False)
+    if pad_widths.min() < 0:
+        raise ValueError("pad_widths must be > 0")
+    pad_widths = np.broadcast_to(pad_widths, (x.ndim, 2)).tolist()
 
     if mode in ['symmetric', 'reflect']:
         xp = np.pad(x, pad_widths, mode=mode)
@@ -504,7 +507,6 @@ def pad_antisymmetric(vector, pad_width, iaxis, kwargs):
             return vector
         xp = np.pad(x, pad_widths, pad_antisymmetric)
     elif mode == 'antireflect':
-        npad_l, npad_r = pad_widths
         xp = np.pad(x, pad_widths, mode='reflect', reflect_type='odd')
     else:
         raise ValueError(

pywt/tests/test_dwt_idwt.py

@@ -2,8 +2,8 @@
 from __future__ import division, print_function, absolute_import
 
 import numpy as np
-from numpy.testing import assert_allclose, assert_, assert_raises
-
+from numpy.testing import (assert_allclose, assert_, assert_raises,
+                           assert_array_equal)
 import pywt
 
 # Check that float32, float64, complex64, complex128 are preserved.
@@ -228,8 +228,72 @@ def test_error_on_continuous_wavelet():
 def test_dwt_zero_size_axes():
     # raise on empty input array
     assert_raises(ValueError, pywt.dwt, [], 'db2')
-  
+
     # >1D case uses a different code path so check there as well
     x = np.ones((1, 4))[0:0, :]  # 2D with a size zero axis
     assert_raises(ValueError, pywt.dwt, x, 'db2', axis=0)
 
+
+def test_pad_1d():
+    x = [1, 2, 3]
+    assert_array_equal(pywt.pad(x, (4, 6), 'periodization'),
+                       [1, 2, 3, 3, 1, 2, 3, 3, 1, 2, 3, 3, 1, 2])
+    assert_array_equal(pywt.pad(x, (4, 6), 'periodic'),
+                       [3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3])
+    assert_array_equal(pywt.pad(x, (4, 6), 'constant'),
+                       [1, 1, 1, 1, 1, 2, 3, 3, 3, 3, 3, 3, 3])
+    assert_array_equal(pywt.pad(x, (4, 6), 'zero'),
+                       [0, 0, 0, 0, 1, 2, 3, 0, 0, 0, 0, 0, 0])
+    assert_array_equal(pywt.pad(x, (4, 6), 'smooth'),
+                       [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+    assert_array_equal(pywt.pad(x, (4, 6), 'symmetric'),
+                       [3, 3, 2, 1, 1, 2, 3, 3, 2, 1, 1, 2, 3])
+    assert_array_equal(pywt.pad(x, (4, 6), 'antisymmetric'),
+                       [3, -3, -2, -1, 1, 2, 3, -3, -2, -1, 1, 2, 3])
+    assert_array_equal(pywt.pad(x, (4, 6), 'reflect'),
+                       [1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2, 1])
+    assert_array_equal(pywt.pad(x, (4, 6), 'antireflect'),
+                       [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+
+    # equivalence of various pad_width formats
+    assert_array_equal(pywt.pad(x, 4, 'periodic'),
+                       pywt.pad(x, (4, 4), 'periodic'))
+
+    assert_array_equal(pywt.pad(x, (4, ), 'periodic'),
+                       pywt.pad(x, (4, 4), 'periodic'))
+
+    assert_array_equal(pywt.pad(x, [(4, 4)], 'periodic'),
+                       pywt.pad(x, (4, 4), 'periodic'))
+
+
+def test_pad_errors():
+    # negative pad width
+    x = [1, 2, 3]
+    assert_raises(ValueError, pywt.pad, x, -2, 'periodic')
+
+    # wrong length pad width
+    assert_raises(ValueError, pywt.pad, x, (1, 1, 1), 'periodic')
+
+    # invalid mode name
+    assert_raises(ValueError, pywt.pad, x, 2, 'bad_mode')
+
+
+def test_pad_nd():
+    for ndim in [2, 3]:
+        x = np.arange(4**ndim).reshape((4, ) * ndim)
+        if ndim == 2:
+            pad_widths = [(2, 1), (2, 3)]
+        else:
+            pad_widths = [(2, 1), ] * ndim
+        for mode in pywt.Modes.modes:
+            xp = pywt.pad(x, pad_widths, mode)
+
+            # expected result is the same as applying along axes separably
+            xp_expected = x.copy()
+            for ax in range(ndim):
+                xp_expected = np.apply_along_axis(pywt.pad,
+                                                  ax,
+                                                  xp_expected,
+                                                  pad_widths=[pad_widths[ax]],
+                                                  mode=mode)
+            assert_array_equal(xp, xp_expected)