Refactor FWT padding (#124)

Author: cthoyt

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "PyWavelets",
     "numpy",
     "torch",
+    "more-itertools",
 ]
 license = {text = "EUPL-1.2"}
 

src/ptwt/_util.py

@@ -12,6 +12,7 @@
 import numpy as np
 import pywt
 import torch
+from more_itertools import grouper
 from typing_extensions import ParamSpec, TypeVar
 
 from .constants import (
@@ -209,7 +210,7 @@ def _get_pad(data_len: int, filt_len: int) -> tuple[int, int]:
     # pad to even singal length.
     padr += data_len % 2
 
-    return padr, padl
+    return padl, padr
 
 
 def _adjust_padding_at_reconstruction(
@@ -789,3 +790,8 @@ def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> RetType:
         return wrapper
 
     return deco
+
+
+def _group_for_symmetric(padding: tuple[int, ...]) -> list[tuple[int, int]]:
+    """Repack the padding tuple for symmetric padding."""
+    return list(reversed(list(grouper(padding, 2))))  # type:ignore[arg-type]

src/ptwt/conv_transform.py

@@ -23,6 +23,7 @@
     _preprocess_coeffs,
     _preprocess_tensor,
     _translate_boundary_strings,
+    _group_for_symmetric,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeff1d
 
@@ -60,13 +61,11 @@ def _fwt_pad(
     pytorch_mode = _translate_boundary_strings(mode)
 
     if padding is None:
-        padr, padl = _get_pad(data.shape[-1], _get_len(wavelet))
-    else:
-        padl, padr = padding
+        padding = _get_pad(data.shape[-1], _get_len(wavelet))
     if pytorch_mode == "symmetric":
-        data_pad = _pad_symmetric(data, [(padl, padr)])
+        data_pad = _pad_symmetric(data, _group_for_symmetric(padding))
     else:
-        data_pad = torch.nn.functional.pad(data, [padl, padr], mode=pytorch_mode)
+        data_pad = torch.nn.functional.pad(data, padding, mode=pytorch_mode)
     return data_pad
 
 

src/ptwt/conv_transform_2.py

@@ -24,6 +24,7 @@
     _preprocess_coeffs,
     _preprocess_tensor,
     _translate_boundary_strings,
+    _group_for_symmetric,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeff2d, WaveletDetailTuple2d
 
@@ -88,16 +89,15 @@ def _fwt_pad2(
     pytorch_mode = _translate_boundary_strings(mode)
 
     if padding is None:
-        padb, padt = _get_pad(data.shape[-2], _get_len(wavelet))
-        padr, padl = _get_pad(data.shape[-1], _get_len(wavelet))
-    else:
-        padl, padr, padt, padb = padding
+        _len_wavelet = _get_len(wavelet)
+        padding = (
+            *_get_pad(data.shape[-1], _len_wavelet),
+            *_get_pad(data.shape[-2], _len_wavelet),
+        )
     if pytorch_mode == "symmetric":
-        data_pad = _pad_symmetric(data, [(padt, padb), (padl, padr)])
+        data_pad = _pad_symmetric(data, _group_for_symmetric(padding))
     else:
-        data_pad = torch.nn.functional.pad(
-            data, [padl, padr, padt, padb], mode=pytorch_mode
-        )
+        data_pad = torch.nn.functional.pad(data, padding, mode=pytorch_mode)
     return data_pad
 
 

src/ptwt/conv_transform_3.py

@@ -24,6 +24,7 @@
     _preprocess_coeffs,
     _preprocess_tensor,
     _translate_boundary_strings,
+    _group_for_symmetric,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeffNd, WaveletDetailDict
 
@@ -91,21 +92,16 @@ def _fwt_pad3(
     pytorch_mode = _translate_boundary_strings(mode)
 
     if padding is None:
-        pad_back, pad_front = _get_pad(data.shape[-3], _get_len(wavelet))
-        pad_bottom, pad_top = _get_pad(data.shape[-2], _get_len(wavelet))
-        pad_right, pad_left = _get_pad(data.shape[-1], _get_len(wavelet))
-    else:
-        pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back = padding
-    if pytorch_mode == "symmetric":
-        data_pad = _pad_symmetric(
-            data, [(pad_front, pad_back), (pad_top, pad_bottom), (pad_left, pad_right)]
+        _len_wavelet = _get_len(wavelet)
+        padding = (
+            *_get_pad(data.shape[-1], _len_wavelet),
+            *_get_pad(data.shape[-2], _len_wavelet),
+            *_get_pad(data.shape[-3], _len_wavelet),
         )
+    if pytorch_mode == "symmetric":
+        data_pad = _pad_symmetric(data, _group_for_symmetric(padding))
     else:
-        data_pad = torch.nn.functional.pad(
-            data,
-            [pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back],
-            mode=pytorch_mode,
-        )
+        data_pad = torch.nn.functional.pad(data, padding, mode=pytorch_mode)
     return data_pad
 
 

tests/test_util.py

@@ -10,6 +10,7 @@
     _fold_axes,
     _pad_symmetric,
     _pad_symmetric_1d,
+    _group_for_symmetric,
     _unfold_axes,
 )
 
@@ -78,3 +79,21 @@ def test_fold(keep_no: int, size: list[int]) -> None:
     assert len(folded.shape) == keep_no + 1
     rec = _unfold_axes(folded, size, keep_no)
     np.allclose(array.numpy(), rec.numpy())
+
+
+def test_repack_symmetric() -> None:
+    """Ensure channel folding works as expected."""
+    padl, padr = padding = tuple(range(2))
+    assert _group_for_symmetric(padding) == [(padl, padr)]
+
+    padl, padr, padt, padb = padding = tuple(range(4))
+    assert _group_for_symmetric(padding) == [(padt, padb), (padl, padr)]
+
+    pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back = padding = tuple(
+        range(6)
+    )
+    assert _group_for_symmetric(padding) == [
+        (pad_front, pad_back),
+        (pad_top, pad_bottom),
+        (pad_left, pad_right),
+    ]