Merge pull request #63 from v0lta/v0.1.6

V0.1.6

Author: v0lta

.flake8

@@ -22,6 +22,7 @@ ignore =
     N400
     # asserts are ok in test.
     S101
+    C901
 exclude =
     .tox,
     .git,

docs/ptwt.rst

@@ -27,14 +27,15 @@ ptwt.conv\_transform\_3 module
    :undoc-members:
    :show-inheritance:
 
-ptwt.separable\_conv\_transform module
---------------------------------------
+ptwt.packets module
+-------------------
 
-.. automodule:: ptwt.separable_conv_transform
+.. automodule:: ptwt.packets
    :members:
    :undoc-members:
    :show-inheritance:
 
+
 ptwt.continuous\_transform module
 ---------------------------------
 
@@ -43,10 +44,11 @@ ptwt.continuous\_transform module
    :undoc-members:
    :show-inheritance:
 
-ptwt.packets module
--------------------
 
-.. automodule:: ptwt.packets
+ptwt.separable\_conv\_transform module
+--------------------------------------
+
+.. automodule:: ptwt.separable_conv_transform
    :members:
    :undoc-members:
    :show-inheritance:

src/ptwt/_stationary_transform.py

@@ -0,0 +1,56 @@
+"""This module implements stationary wavelet transforms."""
+
+from typing import List, Optional, Union
+
+import pywt
+import torch
+
+from src.ptwt._util import Wavelet
+from src.ptwt.conv_transform import _get_filter_tensors
+
+
+def swt(
+    data: torch.Tensor,
+    wavelet: Union[Wavelet, str],
+    level: Optional[int] = None,
+) -> List[torch.Tensor]:
+    """Compute a multilevel 1d stationary wavelet transform.
+
+    Args:
+        data (torch.Tensor): The input data of shape [batch_size, time].
+        wavelet (Union[Wavelet, str]): The wavelet to use.
+        level (Optional[int], optional): The number of levels to compute
+
+    Returns:
+        List[torch.Tensor]: Same as wavedec.
+        Equivalent to pywt.swt with trim_approx=True.
+    """
+    if data.dim() == 1:
+        # assume time series
+        data = data.unsqueeze(0).unsqueeze(0)
+    elif data.dim() == 2:
+        # assume batched time series
+        data = data.unsqueeze(1)
+
+    dec_lo, dec_hi, _, _ = _get_filter_tensors(
+        wavelet, flip=True, device=data.device, dtype=data.dtype
+    )
+    filt_len = dec_lo.shape[-1]
+    filt = torch.stack([dec_lo, dec_hi], 0)
+
+    if level is None:
+        level = pywt.swt_max_level(data.shape[-1])
+
+    result_lst = []
+    res_lo = data
+    for current_level in range(level):
+        dilation = 2**current_level
+        padl, padr = dilation * (filt_len // 2 - 1), dilation * (filt_len // 2)
+        res_lo = torch.nn.functional.pad(res_lo, [padl, padr], mode="circular")
+        res = torch.nn.functional.conv1d(res_lo, filt, stride=1, dilation=dilation)
+        res_lo, res_hi = torch.split(res, 1, 1)
+        # Trim_approx == False
+        # result_lst.append((res_lo.squeeze(1), res_hi.squeeze(1)))
+        result_lst.append(res_hi.squeeze(1))
+    result_lst.append(res_lo.squeeze(1))
+    return result_lst[::-1]

src/ptwt/_util.py

@@ -1,5 +1,5 @@
 """Utility methods to compute wavelet decompositions from a dataset."""
-from typing import Optional, Protocol, Sequence, Tuple, Union
+from typing import List, Optional, Protocol, Sequence, Tuple, Union
 
 import pywt
 import torch
@@ -64,3 +64,57 @@ def _get_len(wavelet: Union[Tuple[torch.Tensor, ...], str, Wavelet]) -> int:
         return wavelet[0].shape[0]
     else:
         return len(_as_wavelet(wavelet))
+
+
+def _pad_symmetric_1d(signal: torch.Tensor, pad_list: Tuple[int, int]) -> torch.Tensor:
+    padl, padr = pad_list
+    dimlen = signal.shape[0]
+    if padl > dimlen or padr > dimlen:
+        if padl > dimlen:
+            signal = _pad_symmetric_1d(signal, (dimlen, 0))
+            padl = padl - dimlen
+        if padr > dimlen:
+            signal = _pad_symmetric_1d(signal, (0, dimlen))
+            padr = padr - dimlen
+        return _pad_symmetric_1d(signal, (padl, padr))
+    else:
+        cat_list = [signal]
+        if padl > 0:
+            topadl = signal[:padl].flip(0)
+            cat_list.insert(0, topadl)
+        if padr > 0:
+            topadr = signal[-padr::].flip(0)
+            cat_list.append(topadr)
+        return torch.cat(cat_list, axis=0)  # type: ignore
+
+
+def _pad_symmetric(
+    signal: torch.Tensor, pad_lists: List[Tuple[int, int]]
+) -> torch.Tensor:
+    if len(signal.shape) < len(pad_lists):
+        raise ValueError("not enough dimensions to pad.")
+
+    dims = len(signal.shape) - 1
+    for pos, pad_list in enumerate(pad_lists[::-1]):
+        current_axis = dims - pos
+        signal = signal.transpose(0, current_axis)
+        signal = _pad_symmetric_1d(signal, pad_list)
+        signal = signal.transpose(current_axis, 0)
+    return signal
+
+
+def _fold_channels(data: torch.Tensor) -> torch.Tensor:
+    """Fold [batch, channel, height width] into [batch*channel, height, widht]."""
+    ds = data.shape
+    fold_data = torch.permute(data, [2, 3, 0, 1])
+    fold_data = torch.reshape(fold_data, [ds[2], ds[3], ds[0] * ds[1]])
+    return torch.permute(fold_data, [2, 0, 1])
+
+
+def _unfold_channels(data: torch.Tensor, ds: List[int]) -> torch.Tensor:
+    """Unfold [batch*channel, height, widht] into [batch, channel, height, width]."""
+    unfold_data = torch.permute(data, [1, 2, 0])
+    unfold_data = torch.reshape(
+        unfold_data, [data.shape[1], data.shape[2], ds[0], ds[1]]
+    )
+    return torch.permute(unfold_data, [2, 3, 0, 1])

src/ptwt/continuous_transform.py

@@ -27,7 +27,7 @@ def cwt(
     wavelet: Union[ContinuousWavelet, str],
     sampling_period: float = 1.0,
 ) -> Tuple[torch.Tensor, np.ndarray]:  # type: ignore
-    """Compute the single dimensional continuous wavelet transform.
+    """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
@@ -166,7 +166,7 @@ def _integrate_wavelet(
 
     Parameters
     ----------
-    wavelet : Wavelet instance or str
+    wavelet: Wavelet instance or str
         Wavelet to integrate.  If a string, should be the name of a wavelet.
     precision : int, optional
         Precision that will be used for wavelet function

src/ptwt/conv_transform.py

@@ -8,7 +8,15 @@
 import pywt
 import torch
 
-from ._util import Wavelet, _as_wavelet, _get_len, _is_dtype_supported
+from ._util import (
+    Wavelet,
+    _as_wavelet,
+    _fold_channels,
+    _get_len,
+    _is_dtype_supported,
+    _pad_symmetric,
+    _unfold_channels,
+)
 
 
 def _create_tensor(
@@ -26,7 +34,7 @@ def _create_tensor(
             return torch.tensor(filter, device=device, dtype=dtype).unsqueeze(0)
 
 
-def get_filter_tensors(
+def _get_filter_tensors(
     wavelet: Union[Wavelet, str],
     flip: bool,
     device: Union[torch.device, str],
@@ -116,6 +124,10 @@ def _translate_boundary_strings(pywt_mode: str) -> str:
         pt_mode = pywt_mode
     elif pywt_mode == "periodic":
         pt_mode = "circular"
+    elif pywt_mode == "symmetric":
+        # pytorch does not support symmetric mode,
+        # we have our own implementation.
+        pt_mode = pywt_mode
     else:
         raise ValueError("Padding mode not supported.")
     return pt_mode
@@ -129,12 +141,12 @@ def _fwt_pad(
     The padding assumes a future step will transform the last axis.
 
     Args:
-        data (torch.Tensor): Input data [batch_size, 1, time]
+        data (torch.Tensor): Input data ``[batch_size, 1, time]``
         wavelet (Wavelet or str): A pywt wavelet compatible object or
             the name of a pywt wavelet.
         mode (str): The desired way to pad. The following methods are supported::
 
-                "reflect", "zero", "constant", "periodic".
+                "reflect", "zero", "constant", "periodic", "symmetric".
 
             Refection padding mirrors samples along the border.
             Zero padding pads zeros.
@@ -151,7 +163,10 @@ def _fwt_pad(
     mode = _translate_boundary_strings(mode)
 
     padr, padl = _get_pad(data.shape[-1], _get_len(wavelet))
-    data_pad = torch.nn.functional.pad(data, [padl, padr], mode=mode)
+    if mode == "symmetric":
+        data_pad = _pad_symmetric(data, [(padl, padr)])
+    else:
+        data_pad = torch.nn.functional.pad(data, [padl, padr], mode=mode)
     return data_pad
 
 
@@ -196,10 +211,40 @@ def _adjust_padding_at_reconstruction(
     elif next_size == pred_size - 1:
         pad_end += 1
     else:
-        raise AssertionError("padding error, please open an issue on github")
+        raise AssertionError(
+            "padding error, please check if dec and rec wavelets are identical."
+        )
     return pad_end, pad_start
 
 
+def _wavedec_fold_channels_1d(data: torch.Tensor) -> Tuple[torch.Tensor, List[int]]:
+    data = data.unsqueeze(-2)
+    ds = data.shape
+    data = _fold_channels(data)
+    return data, list(ds)
+
+
+def _wavedec_unfold_channels_1d_list(
+    result_list: List[torch.Tensor], ds: List[int]
+) -> List[torch.Tensor]:
+    unfold_res = []
+    for res_coeff in result_list:
+        unfold_res.append(
+            _unfold_channels(res_coeff.unsqueeze(1), list(ds)).squeeze(-2)
+        )
+    return unfold_res
+
+
+def _waverec_fold_channels_1d_list(
+    coeff_list: List[torch.Tensor],
+) -> Tuple[List[torch.Tensor], List[int]]:
+    folded = []
+    ds = coeff_list[0].unsqueeze(-2).shape
+    for to_fold_coeff in coeff_list:
+        folded.append(_fold_channels(to_fold_coeff.unsqueeze(-2)).squeeze(-2))
+    return folded, list(ds)
+
+
 def wavedec(
     data: torch.Tensor,
     wavelet: Union[Wavelet, str],
@@ -209,19 +254,27 @@ def wavedec(
     """Compute the analysis (forward) 1d fast wavelet transform.
 
     Args:
-        data (torch.Tensor): Input time series of shape [batch_size, 1, time]
-                             1d inputs are interpreted as [time],
-                             2d inputs are interpreted as [batch_size, time].
+        data (torch.Tensor): The input time series,
+                             1d inputs are interpreted as ``[time]``,
+                             2d inputs as ``[batch_size, time]``,
+                             and 3d inputs as ``[batch_size, channels, time]``.
         wavelet (Wavelet or str): A pywt wavelet compatible object or
             the name of a pywt wavelet.
             Please consider the output from ``pywt.wavelist(kind='discrete')``
             for possible choices.
         mode (str): The desired padding mode. Padding extends the signal along
             the edges. Supported methods are::
 
-                "reflect", "zero", "constant", "periodic".
+                "reflect", "zero", "constant", "periodic", "symmetric".
 
             Defaults to "reflect".
+
+            Symmetric padding mirrors samples along the border.
+            Refection padding reflects samples along the border.
+            Zero padding pads zeros.
+            Constant padding replicates border values.
+            Periodic padding cyclically repeats samples.
+
         level (int): The scale level to be computed.
                                Defaults to None.
 
@@ -246,19 +299,23 @@ def wavedec(
         >>> # compute the forward fwt coefficients
         >>> ptwt.wavedec(data_torch, pywt.Wavelet('haar'),
         >>>              mode='zero', level=2)
-
     """
+    fold = False
     if data.dim() == 1:
         # assume time series
         data = data.unsqueeze(0).unsqueeze(0)
     elif data.dim() == 2:
         # assume batched time series
         data = data.unsqueeze(1)
+    elif data.dim() == 3:
+        # assume batch, channels, time -> fold channels
+        fold = True
+        data, ds = _wavedec_fold_channels_1d(data)
 
     if not _is_dtype_supported(data.dtype):
         raise ValueError(f"Input dtype {data.dtype} not supported")
 
-    dec_lo, dec_hi, _, _ = get_filter_tensors(
+    dec_lo, dec_hi, _, _ = _get_filter_tensors(
         wavelet, flip=True, device=data.device, dtype=data.dtype
     )
     filt_len = dec_lo.shape[-1]
@@ -267,15 +324,20 @@ def wavedec(
     if level is None:
         level = pywt.dwt_max_level(data.shape[-1], filt_len)
 
-    result_lst = []
+    result_list = []
     res_lo = data
     for _ in range(level):
         res_lo = _fwt_pad(res_lo, wavelet, mode=mode)
         res = torch.nn.functional.conv1d(res_lo, filt, stride=2)
         res_lo, res_hi = torch.split(res, 1, 1)
-        result_lst.append(res_hi.squeeze(1))
-    result_lst.append(res_lo.squeeze(1))
-    return result_lst[::-1]
+        result_list.append(res_hi.squeeze(1))
+    result_list.append(res_lo.squeeze(1))
+
+    # unfold if necessary
+    if fold:
+        result_list = _wavedec_unfold_channels_1d_list(result_list, ds)
+
+    return result_list[::-1]
 
 
 def waverec(coeffs: List[torch.Tensor], wavelet: Union[Wavelet, str]) -> torch.Tensor:
@@ -317,7 +379,13 @@ def waverec(coeffs: List[torch.Tensor], wavelet: Union[Wavelet, str]) -> torch.T
         elif torch_dtype != coeff.dtype:
             raise ValueError("coefficients must have the same dtype")
 
-    _, _, rec_lo, rec_hi = get_filter_tensors(
+    # fold channels, if necessary.
+    fold = False
+    if coeffs[0].dim() == 3:
+        fold = True
+        coeffs, ds = _waverec_fold_channels_1d_list(coeffs)
+
+    _, _, rec_lo, rec_hi = _get_filter_tensors(
         wavelet, flip=False, device=torch_device, dtype=torch_dtype
     )
     filt_len = rec_lo.shape[-1]
@@ -339,4 +407,8 @@ def waverec(coeffs: List[torch.Tensor], wavelet: Union[Wavelet, str]) -> torch.T
             res_lo = res_lo[..., padl:]
         if padr > 0:
             res_lo = res_lo[..., :-padr]
+
+    if fold:
+        res_lo = _unfold_channels(res_lo.unsqueeze(-2), list(ds)).squeeze(-2)
+
     return res_lo