Enhance _cwt.py by introducing a configurable hop size parameter

pywt/_cwt.py

@@ -24,9 +24,9 @@ def next_fast_len(n):
     return 2**ceil(np.log2(n))
 
 
-def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
+def cwt(data, scales, wavelet, hop_size=1, sampling_period=1., method='conv', axis=-1):
     """
-    cwt(data, scales, wavelet)
+    cwt(data, scales, wavelet, hop_size)
 
     One dimensional Continuous Wavelet Transform.
 
@@ -41,6 +41,14 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
         ``sampling_period`` is given in seconds.
     wavelet : Wavelet object or name
         Wavelet to use
+    hop_size : int
+        Specifies the down-sampling factor applied on temporal axis during the transform.
+        The output is sampled every hop_size samples, rather than at every consecutive sample.
+        For example:
+        A signal of length 1024 yields 1024 output samples when hop_size=1;
+        512 output samples when hop_size=2;
+        256 output samples when hop_size=4.
+        hop_size must be a positive integer (≥1).
     sampling_period : float
         Sampling period for the frequencies output (optional).
         The values computed for ``coefs`` are independent of the choice of
@@ -73,7 +81,7 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
 
     Notes
     -----
-    Size of coefficients arrays depends on the length of the input array and
+    Size of coefficients arrays depends on the length of the input array, the given hop_size and
     the length of given scales.
 
     Examples
@@ -83,7 +91,7 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
     >>> import matplotlib.pyplot as plt
     >>> x = np.arange(512)
     >>> y = np.sin(2*np.pi*x/32)
-    >>> coef, freqs=pywt.cwt(y,np.arange(1,129),'gaus1')
+    >>> coef, freqs=pywt.cwt(y,np.arange(1,129),1,'gaus1')
     >>> plt.matshow(coef)
     >>> plt.show()
 
@@ -93,7 +101,7 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
     >>> t = np.linspace(-1, 1, 200, endpoint=False)
     >>> sig  = np.cos(2 * np.pi * 7 * t) + np.real(np.exp(-7*(t-0.4)**2)*np.exp(1j*2*np.pi*2*(t-0.4)))
     >>> widths = np.arange(1, 31)
-    >>> cwtmatr, freqs = pywt.cwt(sig, widths, 'mexh')
+    >>> cwtmatr, freqs = pywt.cwt(sig, widths, 2, 'mexh')
     >>> plt.imshow(cwtmatr, extent=[-1, 1, 1, 31], cmap='PRGn', aspect='auto',
     ...            vmax=abs(cwtmatr).max(), vmin=-abs(cwtmatr).max())
     >>> plt.show()
@@ -114,7 +122,8 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
         raise AxisError("axis must be a scalar.")
 
     dt_out = dt_cplx if wavelet.complex_cwt else dt
-    out = np.empty((np.size(scales),) + data.shape, dtype=dt_out)
+    data_sampled = data[..., ::hop_size]
+    out = np.empty((np.size(scales),) + data_sampled.shape, dtype=dt_out)
     precision = 10
     int_psi, x = integrate_wavelet(wavelet, precision=precision)
     int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
@@ -187,7 +196,7 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
             # restore original data shape and axis position
             coef = coef.reshape(data_shape_pre)
             coef = coef.swapaxes(axis, -1)
-        out[i, ...] = coef
+        out[i, ...] = coef[..., ::hop_size]
 
     frequencies = scale2frequency(wavelet, scales, precision)
     if np.isscalar(frequencies):

pywt/tests/test_cwt_wavelets.py

@@ -379,17 +379,18 @@ def test_cwt_complex(dtype, tol, method):
     dt = time[1] - time[0]
     wavelet = 'cmor1.5-1.0'
     scales = np.arange(1, 32)
+    hop_size = 1
 
     # real-valued tranfsorm as a reference
-    [cfs, f] = pywt.cwt(sst, scales, wavelet, dt, method=method)
+    [cfs, f] = pywt.cwt(sst, scales, wavelet, hop_size, dt, method=method)
 
     # verify same precision
     assert_equal(cfs.real.dtype, sst.dtype)
 
     # complex-valued transform equals sum of the transforms of the real
     # and imaginary components
     sst_complex = sst + 1j*sst
-    [cfs_complex, f] = pywt.cwt(sst_complex, scales, wavelet, dt,
+    [cfs_complex, f] = pywt.cwt(sst_complex, scales, wavelet, hop_size, dt,
                                 method=method)
     assert_allclose(cfs + 1j*cfs, cfs_complex, atol=tol, rtol=tol)
     # verify dtype is preserved
@@ -407,12 +408,13 @@ def test_cwt_batch(axis, method):
     dt = time[1] - time[0]
     wavelet = 'cmor1.5-1.0'
     scales = np.arange(1, 32)
+    hop_size = 1
 
     # non-batch transform as reference
-    [cfs1, f] = pywt.cwt(sst1, scales, wavelet, dt, method=method, axis=axis)
+    [cfs1, f] = pywt.cwt(sst1, scales, wavelet, hop_size, dt, method=method, axis=axis)
 
     shape_in = sst.shape
-    [cfs, f] = pywt.cwt(sst, scales, wavelet, dt, method=method, axis=axis)
+    [cfs, f] = pywt.cwt(sst, scales, wavelet, hop_size, dt, method=method, axis=axis)
 
     # shape of input is not modified
     assert_equal(shape_in, sst.shape)