Merge pull request #27 from v0lta/0.1-release-candidate

version 0.1 - release candidate

Author: v0lta

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.19-dev
+current_version = 0.1.0-dev
 commit = True
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(?:-(?P<release>[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?(?:\+(?P<build>[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?

examples/continuous_chirp_analysis/chirp_cwt.png

@@ -0,0 +1,28 @@
+import torch
+import numpy as np
+import src.ptwt as ptwt
+import matplotlib.pyplot as plt
+import scipy.signal as signal
+
+if __name__ == "__main__":
+    t = np.linspace(-2, 2, 800, endpoint=False)
+    sig = signal.chirp(t, f0=1, f1=12, t1=2, method="linear")
+    widths = np.arange(1, 31)
+    cwtmatr_pt, freqs = ptwt.cwt(
+        torch.from_numpy(sig), widths, "mexh", sampling_period=(4 / 800) * np.pi
+    )
+    cwtmatr = cwtmatr_pt.numpy()
+    fig, axs = plt.subplots(2)
+    axs[0].plot(t, sig)
+    axs[0].set_ylabel("magnitude")
+    axs[1].imshow(
+        cwtmatr,
+        cmap="PRGn",
+        aspect="auto",
+        vmax=abs(cwtmatr).max(),
+        vmin=-abs(cwtmatr).max(),
+        extent=[min(t), max(t), min(freqs), max(freqs)],
+    )
+    axs[1].set_xlabel("time")
+    axs[1].set_ylabel("frequency")
+    plt.show()

examples/continuous_chirp_analysis/cwt_chirp_analysis.py

@@ -0,0 +1,7 @@
+#### Cwt chirp analysis
+To run this example, clone the repository and type
+```ipython examples/wavelet_packet_chirp_analysis/cwt_chirp_analysis.py```
+
+The result should look like this:
+
+![alt text](chirp_cwt.png)

examples/continuous_chirp_analysis/readme.md

@@ -18,7 +18,7 @@
 np_lst = []
 for node in nodes:
     np_lst.append(wp[node])
-viz = np.stack(np_lst)
+viz = np.stack(np_lst).squeeze()
 
 fig, axs = plt.subplots(2)
 axs[0].plot(t, w)

examples/wavelet_packet_chirp_analysis/chirp.png

@@ -3,7 +3,7 @@
 ##########################
 [metadata]
 name = ptwt
-version = 0.0.19-dev
+version = 0.1.0-dev
 description = Differentiable and gpu enabled fast wavelet transforms in PyTorch
 long_description = file: README.rst
 long_description_content_type = text/x-rst

examples/wavelet_packet_chirp_analysis/chirp_analysis.py

@@ -1,5 +1,6 @@
 """Differentiable and gpu enabled fast wavelet transforms in PyTorch."""
 from .conv_transform import wavedec, wavedec2, waverec, waverec2
+from .continuous_transform import cwt
 from .matmul_transform import MatrixWavedec, MatrixWaverec
-from .matmul_transform_2d import MatrixWavedec2d, MatrixWaverec2d
+from .matmul_transform_2 import MatrixWavedec2, MatrixWaverec2
 from .packets import WaveletPacket, WaveletPacket2D

setup.cfg

@@ -0,0 +1,122 @@
+"""PyTorch compatible cwt code.
+
+Based on https://github.com/PyWavelets/pywt/blob/master/pywt/_cwt.py
+"""
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+from pywt import ContinuousWavelet, DiscreteContinuousWavelet, Wavelet
+from pywt._functions import integrate_wavelet, scale2frequency
+from torch.fft import fft, ifft
+
+
+def _next_fast_len(n: int) -> int:
+    """Round up size to the nearest power of two.
+
+    Given a number of samples `n`, returns the next power of two
+    following this number to take advantage of FFT speedup.
+    This fallback is less efficient than `scipy.fftpack.next_fast_len`
+    Taken from:
+    https://github.com/PyWavelets/pywt/blob/master/pywt/_cwt.py
+    """
+    return int(2 ** np.ceil(np.log2(n)))
+
+
+def cwt(
+    data: torch.Tensor,
+    scales: Union[np.ndarray, torch.Tensor],  # type: ignore
+    wavelet: Union[ContinuousWavelet, str],
+    sampling_period: float = 1.0,
+) -> Tuple[torch.Tensor, np.ndarray]:  # type: ignore
+    """Compute the single dimensional continuous wavelet transform.
+
+    This function is a PyTorch port of pywt.cwt as found at:
+    https://github.com/PyWavelets/pywt/blob/master/pywt/_cwt.py
+
+    Args:
+        data (torch.Tensor): The input tensor of shape [batch_size, time].
+        scales (torch.Tensor or np.array):
+            The wavelet scales to use. One can use
+            ``f = pywt.scale2frequency(wavelet, scale)/sampling_period`` to determine
+            what physical frequency, ``f``. Here, ``f`` is in hertz when the
+            ``sampling_period`` is given in seconds.
+            wavelet (str or Wavelet of ContinuousWavelet): The wavelet to work with.
+        wavelet (ContinuousWavelet or str): The continuous wavelet to work with.
+        sampling_period (float): Sampling period for the frequencies output (optional).
+            The values computed for ``coefs`` are independent of the choice of
+            ``sampling_period`` (i.e. ``scales`` is not scaled by the sampling
+            period).
+
+    Raises:
+        ValueError: If a scale is too small for the input signal.
+
+    Returns:
+        Tuple[torch.Tensor, np.ndarray]: A tuple with the transformation matrix
+            and frequencies in this order.
+    """
+    # accept array_like input; make a copy to ensure a contiguous array
+    if not isinstance(wavelet, (ContinuousWavelet, Wavelet)):
+        wavelet = DiscreteContinuousWavelet(wavelet)
+    if type(scales) is torch.Tensor:
+        scales = scales.numpy()
+    elif np.isscalar(scales):
+        scales = np.array([scales])
+    # if not np.isscalar(axis):
+    #    raise np.AxisError("axis must be a scalar.")
+
+    precision = 10
+    int_psi, x = integrate_wavelet(wavelet, precision=precision)
+    if type(wavelet) is ContinuousWavelet:
+        int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
+    int_psi = torch.tensor(int_psi, device=data.device)
+
+    # convert int_psi, x to the same precision as the data
+    x = np.asarray(x, dtype=data.cpu().numpy().real.dtype)
+
+    size_scale0 = -1
+    fft_data = None
+
+    out = []
+    for scale in scales:
+        step = x[1] - x[0]
+        j = np.arange(scale * (x[-1] - x[0]) + 1) / (scale * step)
+        j = j.astype(int)  # floor
+        if j[-1] >= len(int_psi):
+            j = np.extract(j < len(int_psi), j)
+        int_psi_scale = int_psi[j].flip(0)
+
+        # The padding is selected for:
+        # - optimal FFT complexity
+        # - to be larger than the two signals length to avoid circular
+        #   convolution
+        size_scale = _next_fast_len(data.shape[-1] + len(int_psi_scale) - 1)
+        if size_scale != size_scale0:
+            # Must recompute fft_data when the padding size changes.
+            fft_data = fft(data, size_scale, dim=-1)
+        size_scale0 = size_scale
+        fft_wav = fft(int_psi_scale, size_scale, dim=-1)
+        conv = ifft(fft_wav * fft_data, dim=-1)
+        conv = conv[..., : data.shape[-1] + len(int_psi_scale) - 1]
+
+        coef = -np.sqrt(scale) * torch.diff(conv, dim=-1)
+
+        # transform axis is always -1
+        d = (coef.shape[-1] - data.shape[-1]) / 2.0
+        if d > 0:
+            coef = coef[..., int(np.floor(d)) : -int(np.ceil(d))]
+        elif d < 0:
+            raise ValueError("Selected scale of {} too small.".format(scale))
+
+        out.append(coef)
+    out_tensor = torch.stack(out)
+    if type(wavelet) is Wavelet:
+        out_tensor = out_tensor.real
+    else:
+        out_tensor = out_tensor if wavelet.complex_cwt else out_tensor.real
+
+    frequencies = scale2frequency(wavelet, scales, precision)
+    if np.isscalar(frequencies):
+        frequencies = np.array([frequencies])
+    frequencies /= sampling_period
+    return out_tensor, frequencies

src/ptwt/__init__.py

@@ -53,7 +53,7 @@ def _create_tensor(filter: Sequence[float]) -> torch.Tensor:
     return dec_lo_tensor, dec_hi_tensor, rec_lo_tensor, rec_hi_tensor
 
 
-def get_pad(data_len: int, filt_len: int) -> Tuple[int, int]:
+def _get_pad(data_len: int, filt_len: int) -> Tuple[int, int]:
     """Compute the required padding.
 
     Args:
@@ -107,12 +107,12 @@ def fwt_pad(
         # convert pywt to pytorch convention.
         mode = "constant"
 
-    padr, padl = get_pad(data.shape[-1], len(wavelet.dec_lo))
+    padr, padl = _get_pad(data.shape[-1], len(wavelet.dec_lo))
     data_pad = torch.nn.functional.pad(data, [padl, padr], mode=mode)
     return data_pad
 
 
-def fwt_pad2d(
+def fwt_pad2(
     data: torch.Tensor, wavelet: Union[Wavelet, str], level: int, mode: str = "reflect"
 ) -> torch.Tensor:
     """Pad data for the 2d FWT.
@@ -129,8 +129,8 @@ def fwt_pad2d(
 
     """
     wavelet = _as_wavelet(wavelet)
-    padb, padt = get_pad(data.shape[-2], len(wavelet.dec_lo))
-    padr, padl = get_pad(data.shape[-1], len(wavelet.dec_lo))
+    padb, padt = _get_pad(data.shape[-2], len(wavelet.dec_lo))
+    padr, padl = _get_pad(data.shape[-1], len(wavelet.dec_lo))
     data_pad = torch.nn.functional.pad(data, [padl, padr, padt, padb], mode=mode)
     return data_pad
 
@@ -254,7 +254,7 @@ def wavedec2(
     ] = []
     res_ll = data
     for s in range(level):
-        res_ll = fwt_pad2d(res_ll, wavelet, level=s, mode=mode)
+        res_ll = fwt_pad2(res_ll, wavelet, level=s, mode=mode)
         res = torch.nn.functional.conv2d(res_ll, dec_filt, stride=2)
         res_ll, res_lh, res_hl, res_hh = torch.split(res, 1, 1)
         result_lst.append((res_lh, res_hl, res_hh))