Refactor convolutional transforms

src/ptwt/_util.py

@@ -850,3 +850,87 @@ def _get_default_axes(n: int) -> tuple[int, ...]:
     if n < 1:
         raise ValueError(f"only natural number dimensions are allowed. given: {n}")
     return tuple(range(-n, 0))
+
+
+def _preprocess_deconstruction(
+    data: torch.Tensor,
+    wavelet: Union[Wavelet, str],
+    *,
+    ndim: int,
+    axes: AxisHint = None,
+) -> tuple[torch.Tensor, list[int], torch.Tensor, torch.Tensor, torch.Tensor]:
+    data, ds = _preprocess_tensor(data, ndim=ndim, axes=axes)
+    dec_lo, dec_hi, _, _ = _get_filter_tensors(
+        wavelet, flip=True, device=data.device, dtype=data.dtype
+    )
+    dec_filt = _construct_nd_filt(dec_lo, dec_hi, n=ndim)
+    return data, ds, dec_lo, dec_hi, dec_filt
+
+
+def _construct_nd_filt(lo: torch.Tensor, hi: torch.Tensor, n: int) -> torch.Tensor:
+    if n == 1:
+        return _construct_1d_filt(lo, hi)
+    elif n == 2:
+        return _construct_2d_filt(lo, hi)
+    elif n == 3:
+        return _construct_3d_filt(lo, hi)
+    else:
+        raise NotImplementedError()
+
+
+def _construct_1d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
+    """Construct one-dimensional filters."""
+    return torch.stack([lo, hi], 0)
+
+
+def _construct_2d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
+    """Construct two-dimensional filters using outer products.
+
+    Args:
+        lo (torch.Tensor): Low-pass input filter.
+        hi (torch.Tensor): High-pass input filter
+
+    Returns:
+        Stacked 2d-filters of dimension
+
+        [2^2, 1, height, width].
+
+        The four filters are ordered ll, lh, hl, hh.
+
+    """
+    ll = _outer(lo, lo)
+    lh = _outer(hi, lo)
+    hl = _outer(lo, hi)
+    hh = _outer(hi, hi)
+    filt = torch.stack([ll, lh, hl, hh], 0)
+    filt = filt.unsqueeze(1)
+    return filt
+
+
+def _construct_3d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
+    """Construct three-dimensional filters using outer products.
+
+    Args:
+        lo (torch.Tensor): Low-pass input filter.
+        hi (torch.Tensor): High-pass input filter
+
+    Returns:
+        Stacked 3d filters of dimension::
+
+        [2^3, 1, length, height, width].
+
+        The four filters are ordered ll, lh, hl, hh.
+    """
+    dim_size = lo.shape[-1]
+    size = [dim_size] * 3
+    lll = _outer(lo, _outer(lo, lo)).reshape(size)
+    llh = _outer(lo, _outer(lo, hi)).reshape(size)
+    lhl = _outer(lo, _outer(hi, lo)).reshape(size)
+    lhh = _outer(lo, _outer(hi, hi)).reshape(size)
+    hll = _outer(hi, _outer(lo, lo)).reshape(size)
+    hlh = _outer(hi, _outer(lo, hi)).reshape(size)
+    hhl = _outer(hi, _outer(hi, lo)).reshape(size)
+    hhh = _outer(hi, _outer(hi, hi)).reshape(size)
+    filt = torch.stack([lll, llh, lhl, lhh, hll, hlh, hhl, hhh], 0)
+    filt = filt.unsqueeze(1)
+    return filt

src/ptwt/conv_transform.py

@@ -22,7 +22,7 @@
     _postprocess_coeffs,
     _postprocess_tensor,
     _preprocess_coeffs,
-    _preprocess_tensor,
+    _preprocess_deconstruction,
     _translate_boundary_strings,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeff1d
@@ -91,7 +91,7 @@ def wavedec(
 
     Args:
         data (torch.Tensor): The input time series to transform.
-            By default the last axis is transformed.
+            By default, the last axis is transformed.
         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')``
@@ -122,22 +122,19 @@ def wavedec(
         >>> # compute the forward fwt coefficients
         >>> ptwt.wavedec(data, 'haar', mode='zero', level=2)
     """
-    data, ds = _preprocess_tensor(data, ndim=1, axes=axis)
-
-    dec_lo, dec_hi, _, _ = _get_filter_tensors(
-        wavelet, flip=True, device=data.device, dtype=data.dtype
+    data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction(
+        data, wavelet, axes=axis, ndim=1
     )
-    filt_len = dec_lo.shape[-1]
-    filt = torch.stack([dec_lo, dec_hi], 0)
 
     if level is None:
+        filt_len = dec_lo.shape[-1]
         level = pywt.dwt_max_level(data.shape[-1], filt_len)
 
     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 = torch.nn.functional.conv1d(res_lo, dec_filt, stride=2)
         res_lo, res_hi = torch.split(res, 1, 1)
         result_list.append(res_hi.squeeze(1))
     result_list.append(res_lo.squeeze(1))

src/ptwt/conv_transform_2.py

@@ -15,46 +15,22 @@
     AxisHint,
     _adjust_padding_at_reconstruction,
     _check_same_device_dtype,
+    _construct_2d_filt,
     _get_filter_tensors,
     _get_padding_n,
     _group_for_symmetric,
-    _outer,
     _pad_symmetric,
     _postprocess_coeffs,
     _postprocess_tensor,
     _preprocess_coeffs,
-    _preprocess_tensor,
+    _preprocess_deconstruction,
     _translate_boundary_strings,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeff2d, WaveletDetailTuple2d
 
 __all__ = ["wavedec2", "waverec2"]
 
 
-def _construct_2d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
-    """Construct two-dimensional filters using outer products.
-
-    Args:
-        lo (torch.Tensor): Low-pass input filter.
-        hi (torch.Tensor): High-pass input filter
-
-    Returns:
-        Stacked 2d-filters of dimension
-
-        [2^2, 1, height, width].
-
-        The four filters are ordered ll, lh, hl, hh.
-
-    """
-    ll = _outer(lo, lo)
-    lh = _outer(hi, lo)
-    hl = _outer(lo, hi)
-    hh = _outer(hi, hi)
-    filt = torch.stack([ll, lh, hl, hh], 0)
-    filt = filt.unsqueeze(1)
-    return filt
-
-
 def _fwt_pad2(
     data: torch.Tensor,
     wavelet: Union[Wavelet, str],
@@ -154,11 +130,9 @@ def wavedec2(
         >>> coefficients = ptwt.wavedec2(data, "haar", level=2, mode="zero")
 
     """
-    data, ds = _preprocess_tensor(data, ndim=2, axes=axes)
-    dec_lo, dec_hi, _, _ = _get_filter_tensors(
-        wavelet, flip=True, device=data.device, dtype=data.dtype
+    data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction(
+        data, wavelet, axes=axes, ndim=2
     )
-    dec_filt = _construct_2d_filt(lo=dec_lo, hi=dec_hi)
 
     if level is None:
         level = pywt.dwtn_max_level([data.shape[-1], data.shape[-2]], wavelet)

src/ptwt/conv_transform_3.py

@@ -15,51 +15,22 @@
     _adjust_padding_at_reconstruction,
     _as_wavelet,
     _check_same_device_dtype,
+    _construct_3d_filt,
     _get_filter_tensors,
     _get_padding_n,
     _group_for_symmetric,
-    _outer,
     _pad_symmetric,
     _postprocess_coeffs,
     _postprocess_tensor,
     _preprocess_coeffs,
-    _preprocess_tensor,
+    _preprocess_deconstruction,
     _translate_boundary_strings,
 )
 from .constants import BoundaryMode, Wavelet, WaveletCoeffNd, WaveletDetailDict
 
 __all__ = ["wavedec3", "waverec3"]
 
 
-def _construct_3d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor:
-    """Construct three-dimensional filters using outer products.
-
-    Args:
-        lo (torch.Tensor): Low-pass input filter.
-        hi (torch.Tensor): High-pass input filter
-
-    Returns:
-        Stacked 3d filters of dimension::
-
-        [2^3, 1, length, height, width].
-
-        The four filters are ordered ll, lh, hl, hh.
-    """
-    dim_size = lo.shape[-1]
-    size = [dim_size] * 3
-    lll = _outer(lo, _outer(lo, lo)).reshape(size)
-    llh = _outer(lo, _outer(lo, hi)).reshape(size)
-    lhl = _outer(lo, _outer(hi, lo)).reshape(size)
-    lhh = _outer(lo, _outer(hi, hi)).reshape(size)
-    hll = _outer(hi, _outer(lo, lo)).reshape(size)
-    hlh = _outer(hi, _outer(lo, hi)).reshape(size)
-    hhl = _outer(hi, _outer(hi, lo)).reshape(size)
-    hhh = _outer(hi, _outer(hi, hi)).reshape(size)
-    filt = torch.stack([lll, llh, lhl, lhh, hll, hlh, hhl, hhh], 0)
-    filt = filt.unsqueeze(1)
-    return filt
-
-
 def _fwt_pad3(
     data: torch.Tensor,
     wavelet: Union[Wavelet, str],
@@ -137,13 +108,10 @@ def wavedec3(
         >>> data = torch.randn(5, 16, 16, 16)
         >>> transformed = ptwt.wavedec3(data, "haar", level=2, mode="reflect")
     """
-    data, ds = _preprocess_tensor(data, ndim=3, axes=axes)
-
     wavelet = _as_wavelet(wavelet)
-    dec_lo, dec_hi, _, _ = _get_filter_tensors(
-        wavelet, flip=True, device=data.device, dtype=data.dtype
+    data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction(
+        data, wavelet, axes=axes, ndim=3
     )
-    dec_filt = _construct_3d_filt(lo=dec_lo, hi=dec_hi)
 
     if level is None:
         level = pywt.dwtn_max_level(

src/ptwt/matmul_transform_2.py

@@ -15,6 +15,7 @@
     AxisHint,
     _as_wavelet,
     _check_same_device_dtype,
+    _construct_2d_filt,
     _deprecated_alias,
     _ensure_axes,
     _get_filter_tensors,
@@ -32,7 +33,7 @@
     WaveletCoeff2d,
     WaveletDetailTuple2d,
 )
-from .conv_transform_2 import _construct_2d_filt, _fwt_pad2
+from .conv_transform_2 import _fwt_pad2
 from .matmul_transform import (
     BaseMatrixWaveDec,
     construct_boundary_a,