torchwavelets 0.1.0__py3-none-any.whl

torchwavelets/dwt/lowlevel.py

@@ -0,0 +1,153 @@
+"""Low-level 1-D analysis/synthesis filter banks operating on 4-D tensors.
+
+All functions operate on ``(N, C, H, W)`` tensors and filter along one spatial
+axis (``dim=2`` for columns, ``dim=3`` for rows). Filters must be 4-D tensors
+of shape ``(1, 1, L, 1)`` (column) or ``(1, 1, 1, L)`` (row); analysis filters
+must already be reversed (see :mod:`torchwavelets.core.filters`).
+
+Everything is built from plain torch ops (``conv2d`` / ``conv_transpose2d`` /
+index-based padding), so gradients flow through all dtypes, autocast works
+natively, and no custom ``autograd.Function`` is required.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from ..core.modes import pad_indices, padding_sizes
+
+__all__ = ["afb1d", "sfb1d"]
+
+
+def _dim_slice(dim: int, sl) -> tuple[slice, ...]:
+    s = [slice(None)] * 4
+    s[dim] = sl
+    return tuple(s)
+
+
+def _cat_last(x: torch.Tensor, dim: int) -> torch.Tensor:
+    """Append a copy of the last sample along ``dim`` (used for odd lengths)."""
+    idx = torch.tensor([x.shape[dim] - 1], device=x.device)
+    return torch.cat((x, x.index_select(dim, idx)), dim=dim)
+
+
+def afb1d(
+    x: torch.Tensor,
+    h0: torch.Tensor,
+    h1: torch.Tensor,
+    mode: str,
+    dim: int,
+) -> torch.Tensor:
+    """1-D analysis filter bank along ``dim`` with decimation by 2.
+
+    Args:
+        x: 4-D input ``(N, C, H, W)``.
+        h0, h1: 4-D analysis filters (reversed), shape ``(1, 1, L, 1)`` if
+            ``dim == 2`` else ``(1, 1, 1, L)``.
+        mode: boundary mode (see :mod:`torchwavelets.core.modes`).
+        dim: spatial dimension to filter along (2 or 3).
+
+    Returns:
+        ``(N, 2C, ...)`` tensor with the lowpass/highpass outputs interleaved
+        along the channel dimension as ``[lo_0, hi_0, lo_1, hi_1, ...]``.
+    """
+    C = x.shape[1]
+    L = h0.numel()
+    stride = (2, 1) if dim == 2 else (1, 2)
+    h = torch.cat([h0, h1] * C, dim=0)
+
+    if mode == "periodization":
+        # Circular transform with the *reversed* analysis filter (conv2d is a
+        # correlation): c[k] = sum_m h[m] * x[(2k + L/2 - m) mod n], which is
+        # PyWavelets' periodization convention (bit-exact vs pywt.dwt).
+        n = x.shape[dim]
+        if n % 2 == 1:
+            x = _cat_last(x, dim)
+            n += 1
+        shift = L // 2 - 1
+        idx = (torch.arange(n + L - 2, device=x.device) - shift) % n
+        xp = x.index_select(dim, idx)
+        return F.conv2d(xp, h, stride=stride, groups=C)
+
+    n = x.shape[dim]
+    before, after = padding_sizes(n, L, mode)
+    if mode == "zero":
+        shape = list(x.shape)
+        shape[dim] = n + before + after
+        padded = torch.zeros(shape, dtype=x.dtype, device=x.device)
+        padded[_dim_slice(dim, slice(before, before + n))] = x
+    elif mode == "smooth":
+        # PyWavelets 'smooth': linear extrapolation from the edges,
+        # left pad (k = before..1): x0 + k*(x0 - x1), right pad (k = 1..after):
+        # x[-1] + k*(x[-1] - x[-2]).
+        x0 = x[_dim_slice(dim, slice(0, 1))]
+        x1 = x[_dim_slice(dim, slice(1, 2))] if n >= 2 else torch.zeros_like(x0)
+        xm1 = x[_dim_slice(dim, slice(-1, None))]
+        xm2 = x[_dim_slice(dim, slice(-2, -1))] if n >= 2 else torch.zeros_like(xm1)
+        kshape = [1, 1, 1, 1]
+        kshape[dim] = -1
+        k_left = torch.arange(before, 0, -1, dtype=x.dtype, device=x.device).view(kshape)
+        k_right = torch.arange(1, after + 1, dtype=x.dtype, device=x.device).view(kshape)
+        left = x0 + k_left * (x0 - x1)
+        right = xm1 + k_right * (xm1 - xm2)
+        padded = torch.cat([left, x, right], dim=dim)
+    else:
+        idx = pad_indices(n, before, after, mode, x.device)
+        padded = x.index_select(dim, idx)
+    return F.conv2d(padded, h, stride=stride, groups=C)
+
+
+def sfb1d(
+    lo: torch.Tensor,
+    hi: torch.Tensor,
+    g0: torch.Tensor,
+    g1: torch.Tensor,
+    mode: str,
+    dim: int,
+) -> torch.Tensor:
+    """1-D synthesis filter bank along ``dim`` with upsampling by 2.
+
+    Args:
+        lo, hi: 4-D lowpass/highpass coefficients ``(N, C, ...)``.
+        g0, g1: 4-D synthesis filters (natural order), shape ``(1, 1, L, 1)``
+            if ``dim == 2`` else ``(1, 1, 1, L)``.
+        mode: boundary mode.
+        dim: spatial dimension (2 or 3).
+
+    Returns:
+        Reconstructed ``(N, C, ...)`` tensor.
+    """
+    C = lo.shape[1]
+    L = g0.numel()
+    stride = (2, 1) if dim == 2 else (1, 2)
+    g0 = torch.cat([g0] * C, dim=0)
+    g1 = torch.cat([g1] * C, dim=0)
+
+    if mode == "periodization":
+        # Circular synthesis: y[m] = sum_j g[j] * u[(m + j - L/2) mod 2n] with
+        # u the upsampled (zero-inserted) coefficient signal. g is flipped
+        # because conv2d is a correlation (PyWavelets' periodization, exact).
+        n_in = lo.shape[dim]
+        n_out = 2 * n_in
+        u_lo = torch.zeros(
+            *lo.shape[:dim], n_out, *lo.shape[dim + 1 :], dtype=lo.dtype, device=lo.device
+        )
+        u_hi = torch.zeros_like(u_lo)
+        u_lo[_dim_slice(dim, slice(None, None, 2))] = lo
+        u_hi[_dim_slice(dim, slice(None, None, 2))] = hi
+        idx = (torch.arange(n_out + L - 1, device=lo.device) - L // 2) % n_out
+        # Flip along the filter-tap dimension: conv2d is a correlation.
+        # g0/g1 are already tiled to C channels above.
+        flip = -2 if dim == 2 else -1
+        g0 = g0.flip(flip)
+        g1 = g1.flip(flip)
+        y = F.conv2d(u_lo.index_select(dim, idx), g0, groups=C) + F.conv2d(
+            u_hi.index_select(dim, idx), g1, groups=C
+        )
+        return y
+
+    pad = (L - 2, 0) if dim == 2 else (0, L - 2)
+    y = F.conv_transpose2d(lo, g0, stride=stride, padding=pad, groups=C)
+    y = y + F.conv_transpose2d(hi, g1, stride=stride, padding=pad, groups=C)
+    return y

torchwavelets/dwt/transform1d.py

@@ -0,0 +1,175 @@
+"""1-D discrete wavelet transform (functional and module API).
+
+The module API resolves its filters into registered buffers at construction
+time (eager) and only performs torch ops (``.to(dtype, device)`` + convs) in
+``forward``, so it is fully traceable by ``torch.compile``.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..core.filters import filter_bank, resolve_1d, wave_name
+from ..core.modes import validate_mode
+from . import lowlevel
+from ._validation import check_reflect_length, check_spatial, validate_level, warn_if_too_high
+
+__all__ = ["dwt1d", "idwt1d", "DWT1D", "IDWT1D"]
+
+_IS_COMPILING = getattr(torch.compiler, "is_compiling", lambda: False)
+
+
+def _dwt1d_impl(x: torch.Tensor, h0, h1, mode: str, level: int):
+    yl = x[:, :, None, :]
+    highs = []
+    for _ in range(level):
+        lohi = lowlevel.afb1d(yl, h0, h1, mode, 3)
+        yl = lohi[:, ::2]
+        highs.append(lohi[:, 1::2][:, :, 0])
+    return yl[:, :, 0], highs
+
+
+def _idwt1d_impl(yl: torch.Tensor, highs, g0, g1, mode: str) -> torch.Tensor:
+    y = yl[:, :, None, :]
+    for h in reversed(highs):
+        if h is None:
+            h = torch.zeros_like(y[:, :, 0])
+        elif y.shape[3] == h.shape[2] + 1:
+            # PyWavelets convention: drop the extra sample introduced by an
+            # odd-length previous level before synthesizing.
+            y = y[..., :-1]
+        if y.shape[3] != h.shape[2]:
+            raise ValueError(
+                "coefficient shape mismatch: approximation "
+                f"{tuple(y.shape)} vs detail {tuple(h.shape)} at the same level"
+            )
+        y = lowlevel.sfb1d(y, h[:, :, None, :], g0, g1, mode, 3)
+    return y[:, :, 0]
+
+
+def dwt1d(
+    x: torch.Tensor, wave="db1", mode: str = "symmetric", level: int = 1
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
+    """Forward 1-D DWT.
+
+    Args:
+        x (torch.Tensor): ``(N, C, L)`` input.
+        wave (str or pywt.Wavelet or tuple): wavelet name, ``pywt.Wavelet`` or
+            custom filter bank.
+        mode (str): boundary mode.
+        level (int): number of decomposition levels (``0`` returns ``(x, [])``).
+
+    Returns:
+        tuple[torch.Tensor, list[torch.Tensor]]: ``(yl, yh)`` where ``yl`` is
+        the final lowpass of shape ``(N, C, L')`` and ``yh`` is a list of
+        ``level`` highpass tensors, finest first. The output dtype matches the
+        input dtype.
+    """
+    if x.ndim != 3:
+        raise ValueError(f"dwt1d expects a 3-D input (N, C, L); got shape {tuple(x.shape)}")
+    check_spatial(x.shape[-1:], ("L",), "dwt1d")
+    validate_level(level)
+    mode = validate_mode(mode)
+    check_reflect_length(x.shape[-1:], ("L",), mode, "dwt1d")
+
+    warn_if_too_high(x.shape[-1], filter_bank(wave).length, level)
+    h0, h1, _, _ = resolve_1d(wave, x.dtype, x.device)
+    return _dwt1d_impl(x, h0, h1, mode, level)
+
+
+def idwt1d(coeffs, wave="db1", mode: str = "symmetric") -> torch.Tensor:
+    """Inverse 1-D DWT.
+
+    Args:
+        coeffs (tuple): ``(yl, yh)`` as returned by :func:`dwt1d` (any
+            highpass entry may be ``None``, in which case it is treated as
+            zeros).
+        wave (str or pywt.Wavelet or tuple): as in :func:`dwt1d`.
+        mode (str): as in :func:`dwt1d`.
+
+    Returns:
+        torch.Tensor: reconstructed ``(N, C, L')`` tensor (dtype matches the
+        coefficients).
+
+    This is a ``pywt.waverec``-style inverse: the coefficient pyramid from
+    :func:`dwt1d` is inverted level by level, and (like ``pywt.waverec``) an
+    odd-length previous level's extra sample is dropped inside the
+    reconstruction. Unlike ``pywt.idwt`` (single-level), the lowpass/detail
+    pair of a level does not need to match in length.
+    """
+    yl, highs = coeffs
+    if yl is None:
+        raise ValueError("The lowpass coefficient tensor yl cannot be None")
+    if yl.ndim != 3:
+        raise ValueError(f"idwt1d expects 3-D coefficients; got shape {tuple(yl.shape)}")
+    check_spatial(yl.shape[-1:], ("L",), "idwt1d")
+    mode = validate_mode(mode)
+
+    _, _, g0, g1 = resolve_1d(wave, yl.dtype, yl.device)
+    return _idwt1d_impl(yl, highs, g0, g1, mode)
+
+
+class DWT1D(torch.nn.Module):
+    """1-D DWT module.
+
+    Args:
+        J (int): number of decomposition levels.
+        wave (str or pywt.Wavelet or tuple): as in :func:`dwt1d`.
+        mode (str): as in :func:`dwt1d`.
+
+    The precision follows the input tensor: feed a float16 tensor and the
+    output is float16, regardless of the dtype used at construction time.
+    The module is ``torch.compile``-compatible (filters are resolved eagerly
+    at construction; only torch ops run in ``forward``).
+    """
+
+    def __init__(self, J: int = 1, wave="db1", mode: str = "symmetric"):
+        super().__init__()
+        validate_level(J)
+        self.J = J
+        self.wave = wave
+        self.mode = validate_mode(mode)
+        self._bank_length = filter_bank(wave).length
+        h0, h1, _, _ = resolve_1d(wave, torch.float64, torch.device("cpu"))
+        self.register_buffer("h0", h0)
+        self.register_buffer("h1", h1)
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim != 3:
+            raise ValueError(f"DWT1D expects a 3-D input (N, C, L); got shape {tuple(x.shape)}")
+        check_spatial(x.shape[-1:], ("L",), "DWT1D")
+        check_reflect_length(x.shape[-1:], ("L",), self.mode, "DWT1D")
+        if not _IS_COMPILING():
+            warn_if_too_high(x.shape[-1], self._bank_length, self.J)
+        h0 = self.h0.to(dtype=x.dtype, device=x.device)
+        h1 = self.h1.to(dtype=x.dtype, device=x.device)
+        return _dwt1d_impl(x, h0, h1, self.mode, self.J)
+
+    def extra_repr(self) -> str:
+        return f"J={self.J}, wave={wave_name(self.wave)!r}, mode={self.mode!r}"
+
+
+class IDWT1D(torch.nn.Module):
+    """1-D inverse DWT module (see :func:`idwt1d`)."""
+
+    def __init__(self, wave="db1", mode: str = "symmetric"):
+        super().__init__()
+        self.wave = wave
+        self.mode = validate_mode(mode)
+        _, _, g0, g1 = resolve_1d(wave, torch.float64, torch.device("cpu"))
+        self.register_buffer("g0", g0)
+        self.register_buffer("g1", g1)
+
+    def forward(self, coeffs):
+        yl, _ = coeffs
+        if yl is None:
+            raise ValueError("The lowpass coefficient tensor yl cannot be None")
+        if yl.ndim != 3:
+            raise ValueError(f"IDWT1D expects 3-D coefficients; got shape {tuple(yl.shape)}")
+        check_spatial(yl.shape[-1:], ("L",), "IDWT1D")
+        g0 = self.g0.to(dtype=yl.dtype, device=yl.device)
+        g1 = self.g1.to(dtype=yl.dtype, device=yl.device)
+        return _idwt1d_impl(yl, coeffs[1], g0, g1, self.mode)
+
+    def extra_repr(self) -> str:
+        return f"wave={wave_name(self.wave)!r}, mode={self.mode!r}"

torchwavelets/dwt/transform2d.py

@@ -0,0 +1,195 @@
+"""2-D discrete wavelet transform (functional and module API).
+
+The module API resolves its filters into registered buffers at construction
+time (eager) and only performs torch ops (``.to(dtype, device)`` + convs) in
+``forward``, so it is fully traceable by ``torch.compile``.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..core.filters import filter_bank, resolve_2d, wave_name
+from ..core.modes import validate_mode
+from . import lowlevel
+from ._validation import check_reflect_length, check_spatial, validate_level, warn_if_too_high
+
+__all__ = ["dwt2d", "idwt2d", "DWT2D", "IDWT2D"]
+
+_IS_COMPILING = getattr(torch.compiler, "is_compiling", lambda: False)
+
+_DWT2D_NAMES = ("h0_col", "h1_col", "h0_row", "h1_row")
+_IDWT2D_NAMES = ("g0_col", "g1_col", "g0_row", "g1_row")
+
+
+def _dwt2d_impl(x: torch.Tensor, h0c, h1c, h0r, h1r, mode: str, level: int):
+    yl = x
+    highs = []
+    for _ in range(level):
+        lohi = lowlevel.afb1d(yl, h0r, h1r, mode, 3)  # rows
+        y = lowlevel.afb1d(lohi, h0c, h1c, mode, 2)  # columns
+        s = y.shape
+        y = y.reshape(s[0], s[1] // 4, 4, s[-2], s[-1])
+        yl = y[:, :, 0].contiguous()
+        highs.append(y[:, :, 1:].contiguous())
+    return yl, highs
+
+
+def _idwt2d_impl(yl: torch.Tensor, highs, g0c, g1c, g0r, g1r, mode: str) -> torch.Tensor:
+    y = yl
+    for h in reversed(highs):
+        if h is None:
+            h = torch.zeros(
+                y.shape[0],
+                y.shape[1],
+                3,
+                y.shape[-2],
+                y.shape[-1],
+                dtype=y.dtype,
+                device=y.device,
+            )
+        else:
+            # PyWavelets convention: drop extra samples introduced by
+            # odd-length previous levels before synthesizing.
+            if y.shape[-2] == h.shape[-2] + 1:
+                y = y[..., :-1, :]
+            if y.shape[-1] == h.shape[-1] + 1:
+                y = y[..., :-1]
+        if y.shape[-2] != h.shape[-2] or y.shape[-1] != h.shape[-1]:
+            raise ValueError(
+                "coefficient shape mismatch: approximation "
+                f"{tuple(y.shape)} vs detail {tuple(h.shape)} at the same level"
+            )
+        lh, hl, hh = torch.unbind(h, dim=2)
+        lo = lowlevel.sfb1d(y, lh, g0c, g1c, mode, 2)
+        hi = lowlevel.sfb1d(hl, hh, g0c, g1c, mode, 2)
+        y = lowlevel.sfb1d(lo, hi, g0r, g1r, mode, 3)
+    return y
+
+
+def dwt2d(
+    x: torch.Tensor, wave="db1", mode: str = "symmetric", level: int = 1
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
+    """Forward 2-D DWT (separable, row then column filtering).
+
+    Args:
+        x (torch.Tensor): ``(N, C, H, W)`` input.
+        wave (str or pywt.Wavelet or tuple): wavelet name, ``pywt.Wavelet`` or
+            custom filter bank.
+        mode (str): boundary mode.
+        level (int): number of decomposition levels (``0`` returns ``(x, [])``).
+
+    Returns:
+        ``(yl, yh)`` where ``yl`` is the final lowpass ``(N, C, H', W')`` and
+        ``yh`` is a list of ``level`` tensors of shape
+        ``(N, C, 3, H'', W'')`` holding the ``(LH, HL, HH)`` subbands
+        (finest first). ``LH``/``HL``/``HH`` match PyWavelets'
+        ``cH``/``cV``/``cD``. The output dtype matches the input dtype.
+    """
+    if x.ndim != 4:
+        raise ValueError(f"dwt2d expects a 4-D input (N, C, H, W); got shape {tuple(x.shape)}")
+    check_spatial(x.shape[2:], ("H", "W"), "dwt2d")
+    mode = validate_mode(mode)
+    check_reflect_length(x.shape[2:], ("H", "W"), mode, "dwt2d")
+    validate_level(level)
+    mode = validate_mode(mode)
+
+    warn_if_too_high(min(x.shape[-2:]), filter_bank(wave).length, level)
+    h0c, h1c, h0r, h1r, _, _, _, _ = resolve_2d(wave, x.dtype, x.device)
+    return _dwt2d_impl(x, h0c, h1c, h0r, h1r, mode, level)
+
+
+def idwt2d(coeffs, wave="db1", mode: str = "symmetric") -> torch.Tensor:
+    """Inverse 2-D DWT.
+
+    Args:
+        coeffs (tuple): ``(yl, yh)`` as returned by :func:`dwt2d` (any
+            highpass entry may be ``None``, in which case it is treated as
+            zeros).
+        wave (str or pywt.Wavelet or tuple): as in :func:`dwt2d`.
+        mode (str): as in :func:`dwt2d`.
+
+    Returns:
+        torch.Tensor: reconstructed ``(N, C, H', W')`` tensor (dtype matches
+        the coefficients).
+    """
+    yl, highs = coeffs
+    if yl is None:
+        raise ValueError("The lowpass coefficient tensor yl cannot be None")
+    if yl.ndim != 4:
+        raise ValueError(f"idwt2d expects 4-D coefficients; got shape {tuple(yl.shape)}")
+    check_spatial(yl.shape[2:], ("H", "W"), "idwt2d")
+    mode = validate_mode(mode)
+
+    _, _, _, _, g0c, g1c, g0r, g1r = resolve_2d(wave, yl.dtype, yl.device)
+    return _idwt2d_impl(yl, highs, g0c, g1c, g0r, g1r, validate_mode(mode))
+
+
+class DWT2D(torch.nn.Module):
+    """2-D DWT module.
+
+    Args:
+        J (int): number of decomposition levels.
+        wave (str or pywt.Wavelet or tuple): as in :func:`dwt2d`.
+        mode (str): as in :func:`dwt2d`.
+
+    The precision follows the input tensor: feed a float16 tensor and the
+    output is float16, regardless of the dtype used at construction time.
+    The module is ``torch.compile``-compatible (filters are resolved eagerly
+    at construction; only torch ops run in ``forward``).
+    """
+
+    def __init__(self, J: int = 1, wave="db1", mode: str = "symmetric"):
+        super().__init__()
+        validate_level(J)
+        self.J = J
+        self.wave = wave
+        self.mode = validate_mode(mode)
+        self._bank_length = filter_bank(wave).length
+        filts = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        for name, t in zip(_DWT2D_NAMES, filts[:4], strict=True):
+            self.register_buffer(name, t)
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim != 4:
+            raise ValueError(f"DWT2D expects a 4-D input (N, C, H, W); got shape {tuple(x.shape)}")
+        check_spatial(x.shape[2:], ("H", "W"), "DWT2D")
+        check_reflect_length(x.shape[2:], ("H", "W"), self.mode, "DWT2D")
+        if not _IS_COMPILING():
+            warn_if_too_high(min(x.shape[-2:]), self._bank_length, self.J)
+        h0c = self.h0_col.to(dtype=x.dtype, device=x.device)
+        h1c = self.h1_col.to(dtype=x.dtype, device=x.device)
+        h0r = self.h0_row.to(dtype=x.dtype, device=x.device)
+        h1r = self.h1_row.to(dtype=x.dtype, device=x.device)
+        return _dwt2d_impl(x, h0c, h1c, h0r, h1r, self.mode, self.J)
+
+    def extra_repr(self) -> str:
+        return f"J={self.J}, wave={wave_name(self.wave)!r}, mode={self.mode!r}"
+
+
+class IDWT2D(torch.nn.Module):
+    """2-D inverse DWT module (see :func:`idwt2d`)."""
+
+    def __init__(self, wave="db1", mode: str = "symmetric"):
+        super().__init__()
+        self.wave = wave
+        self.mode = validate_mode(mode)
+        filts = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        for name, t in zip(_IDWT2D_NAMES, filts[4:], strict=True):
+            self.register_buffer(name, t)
+
+    def forward(self, coeffs):
+        yl, _ = coeffs
+        if yl is None:
+            raise ValueError("The lowpass coefficient tensor yl cannot be None")
+        if yl.ndim != 4:
+            raise ValueError(f"IDWT2D expects 4-D coefficients; got shape {tuple(yl.shape)}")
+        check_spatial(yl.shape[2:], ("H", "W"), "IDWT2D")
+        g0c = self.g0_col.to(dtype=yl.dtype, device=yl.device)
+        g1c = self.g1_col.to(dtype=yl.dtype, device=yl.device)
+        g0r = self.g0_row.to(dtype=yl.dtype, device=yl.device)
+        g1r = self.g1_row.to(dtype=yl.dtype, device=yl.device)
+        return _idwt2d_impl(yl, coeffs[1], g0c, g1c, g0r, g1r, self.mode)
+
+    def extra_repr(self) -> str:
+        return f"wave={wave_name(self.wave)!r}, mode={self.mode!r}"