Merge pull request #506 from grlee77/guard_against_size0_axes

Guard against trying to transform along size 0 axes

Author: rgommers

pywt/_extensions/_dwt.pyx

@@ -90,7 +90,7 @@ cpdef dwt_axis(np.ndarray data, Wavelet wavelet, MODE mode, unsigned int axis=0)
 
     input_shape = <size_t [:data.ndim]> <size_t *> data.shape
     output_shape = input_shape.copy()
-    output_shape[axis] = common.dwt_buffer_length(data.shape[axis], wavelet.dec_len, mode)
+    output_shape[axis] = dwt_coeff_len(data.shape[axis], wavelet.dec_len, mode)
 
     cA = np.empty(output_shape, data.dtype)
     cD = np.empty(output_shape, data.dtype)

pywt/_extensions/_swt.pyx

@@ -36,6 +36,8 @@ def swt_max_level(size_t input_len):
     multiple of ``2**n``. ``numpy.pad`` can be used to pad a signal up to an
     appropriate length as needed.
     """
+    if input_len < 1:
+        raise ValueError("Cannot apply swt to a size 0 signal.")
     max_level = common.swt_max_level(input_len)
     if max_level == 0:
         warnings.warn(
@@ -54,6 +56,8 @@ def swt(cdata_t[::1] data, Wavelet wavelet, size_t level, size_t start_level):
 
     if data.size % 2:
         raise ValueError("Length of data must be even.")
+    if data.size < 1:
+        raise ValueError("Data must have non-zero size")
 
     if level < 1:
         raise ValueError("Level value must be greater than zero.")
@@ -67,6 +71,7 @@ def swt(cdata_t[::1] data, Wavelet wavelet, size_t level, size_t start_level):
                 common.swt_max_level(data.size) - start_level))
         raise ValueError(msg)
 
+
     output_len = common.swt_buffer_length(data.size)
     if output_len < 1:
         raise RuntimeError("Invalid output length.")
@@ -153,8 +158,10 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
     cdef int retval = -5
     cdef size_t i
 
-    if data.size % 2:
-        raise ValueError("Length of data must be even.")
+    if data.shape[axis] % 2:
+        raise ValueError("Length of data must be even along the transform axis.")
+    if data.shape[axis] < 1:
+        raise ValueError("Data must have non-zero size along the transform axis.")
 
     if level < 1:
         raise ValueError("Level value must be greater than zero.")

pywt/tests/test_dwt_idwt.py

@@ -223,3 +223,13 @@ def test_error_on_continuous_wavelet():
 
         cA, cD = pywt.dwt(data, 'db1')
         assert_raises(ValueError, pywt.idwt, cA, cD, cwave)
+
+
+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)
+

pywt/tests/test_swt.py

@@ -525,3 +525,13 @@ def test_iswtn_mixed_dtypes():
         y = pywt.iswtn(coeffs, wav)
         assert_equal(output_dtype, y.dtype)
         assert_allclose(y, x, rtol=1e-3, atol=1e-3)
+
+
+def test_swt_zero_size_axes():
+    # raise on empty input array
+    assert_raises(ValueError, pywt.swt, [], '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.swtn, x, 'db2', level=1, axes=(0,))
+