chuchichaestli.dwt.functional

Separable N-dimensional discrete wavelet transforms.

Subbands carry one character per transformed axis, 'a' for the approximation and 'd' for the detail branch, so a two-dimensional transform yields aa, ad, da and dd (the classical LL, LH, HL, HH).

Functions:

Name	Description
dwt	Single-level wavelet transform along one axis.
dwt_coeff_len	Number of coefficients one level of the transform produces per subband.
dwt_max_level	Number of levels a signal supports before it is shorter than the filter.
dwtn	Single-level separable wavelet transform over one or more axes.
dwtn_approx	Approximation band of a single-level transform.
idwt	Invert a single-level wavelet transform along one axis.
idwtn	Invert a single-level separable wavelet transform.
subband_keys	Subband names of a dimensions-dimensional transform, in output order.
wavedec	Multi-level wavelet transform along one axis.
wavedecn	Multi-level separable wavelet transform, recursing on the approximation.
waverec	Invert a multi-level wavelet transform along one axis.
waverecn	Invert a multi-level separable wavelet transform.

dwt

dwt(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
) -> tuple[torch.Tensor, torch.Tensor]

Single-level wavelet transform along one axis.

Parameters:

Name	Type	Description	Default
data	torch.Tensor	Input tensor of any rank.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode.	'zero'
axis	int	Axis to transform.	-1

Returns:

Type	Description
tuple[torch.Tensor, torch.Tensor]	The approximation and the detail coefficients.

Source code in src/chuchichaestli/dwt/functional.py

def dwt(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Single-level wavelet transform along one axis.

    Args:
        data: Input tensor of any rank.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode.
        axis: Axis to transform.

    Returns:
        The approximation and the detail coefficients.
    """
    bands = dwtn(data, wavelet, mode, (axis,))
    return bands["a"], bands["d"]

dwt_coeff_len

dwt_coeff_len(
    data_len: int,
    filter_len: int,
    mode: ExtensionModeTypes = "zero",
) -> int

Number of coefficients one level of the transform produces per subband.

Parameters:

Name	Type	Description	Default
data_len	int	Length of the signal.	required
filter_len	int	Length of the decomposition filters.	required
mode	ExtensionModeTypes	Signal extension mode.	'zero'

Raises:

Type	Description
ValueError	If data_len or filter_len is not positive.

Source code in src/chuchichaestli/dwt/functional.py

def dwt_coeff_len(
    data_len: int, filter_len: int, mode: ExtensionModeTypes = "zero"
) -> int:
    """Number of coefficients one level of the transform produces per subband.

    Args:
        data_len: Length of the signal.
        filter_len: Length of the decomposition filters.
        mode: Signal extension mode.

    Raises:
        ValueError: If `data_len` or `filter_len` is not positive.
    """
    if data_len < 1 or filter_len < 1:
        raise ValueError(
            f"Signal and filter must be non-empty; got {data_len} and {filter_len}."
        )
    if mode == "periodization":
        return (data_len + 1) // 2
    return (data_len + filter_len - 1) // 2

dwt_max_level

dwt_max_level(data_len: int, filter_len: int) -> int

Number of levels a signal supports before it is shorter than the filter.

Parameters:

Name	Type	Description	Default
data_len	int	Length of the signal.	required
filter_len	int	Length of the decomposition filters.	required

Source code in src/chuchichaestli/dwt/functional.py

def dwt_max_level(data_len: int, filter_len: int) -> int:
    """Number of levels a signal supports before it is shorter than the filter.

    Args:
        data_len: Length of the signal.
        filter_len: Length of the decomposition filters.
    """
    if filter_len < 2 or data_len < filter_len - 1:
        return 0
    return int(torch.log2(torch.tensor(data_len / (filter_len - 1.0))).item())

dwtn

dwtn(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
) -> dict[str, torch.Tensor]

Single-level separable wavelet transform over one or more axes.

Parameters:

Name	Type	Description	Default
data	torch.Tensor	Input tensor of any rank.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode.	'zero'
axes	Sequence[int] | None	Axes to transform; the trailing one if omitted.	None

Returns:

Type	Description
dict[str, torch.Tensor]	The subbands, keyed by their 'a'/'d' names.

Source code in src/chuchichaestli/dwt/functional.py

def dwtn(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
) -> dict[str, torch.Tensor]:
    """Single-level separable wavelet transform over one or more axes.

    Args:
        data: Input tensor of any rank.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode.
        axes: Axes to transform; the trailing one if omitted.

    Returns:
        The subbands, keyed by their `'a'`/`'d'` names.
    """
    data = as_inexact(data)
    axes = _resolve_axes(data.ndim, axes)
    wavelet = as_wavelet(wavelet)
    dec_lo, dec_hi, _, _ = wavelet.filters(data.dtype, data.device)
    from chuchichaestli.dwt import _ext

    h, lead, perm = _fold(data, axes)
    spatial = tuple(h.shape[2:])
    if _ext.kernels_available(h.device, h.dtype) and _ext.fused_haar_applies(
        wavelet, mode, spatial
    ):
        h = _ext.haar_nd(h, len(axes))
    else:
        h = _decompose_axes(h, dec_lo, dec_hi, len(axes), mode)
    keys = subband_keys(len(axes))
    return {key: _unfold(h[:, i], lead, perm) for i, key in enumerate(keys)}

dwtn_approx

dwtn_approx(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
) -> torch.Tensor

Approximation band of a single-level transform.

The compiled kernel splits every band at once for less than a convolution costs to pad for, so it is taken where it serves and its detail bands dropped; the convolution only skips them where there is no kernel.

Parameters:

Name	Type	Description	Default
data	torch.Tensor	Input tensor of any rank.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode.	'zero'
axes	Sequence[int] | None	Axes to transform; the trailing one if omitted.	None

Source code in src/chuchichaestli/dwt/functional.py

def dwtn_approx(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
) -> torch.Tensor:
    """Approximation band of a single-level transform.

    The compiled kernel splits every band at once for less than a convolution
    costs to pad for, so it is taken where it serves and its detail bands
    dropped; the convolution only skips them where there is no kernel.

    Args:
        data: Input tensor of any rank.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode.
        axes: Axes to transform; the trailing one if omitted.
    """
    data = as_inexact(data)
    axes = _resolve_axes(data.ndim, axes)
    dec_lo, _, _, _ = as_wavelet(wavelet).filters(data.dtype, data.device)
    h, lead, perm = _fold(data, axes)
    for axis in range(len(axes)):
        h = _decompose_lowpass(h, dec_lo, axis, mode)
    return _unfold(h[:, 0], lead, perm)

idwt

idwt(
    approx: torch.Tensor,
    detail: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
    output_size: int | None = None,
) -> torch.Tensor

Invert a single-level wavelet transform along one axis.

Parameters:

Name	Type	Description	Default
approx	torch.Tensor	Approximation coefficients.	required
detail	torch.Tensor	Detail coefficients.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode the decomposition used.	'zero'
axis	int	Axis that was transformed.	-1
output_size	int | None	Length of the reconstructed axis.	None

Source code in src/chuchichaestli/dwt/functional.py

def idwt(
    approx: torch.Tensor,
    detail: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
    output_size: int | None = None,
) -> torch.Tensor:
    """Invert a single-level wavelet transform along one axis.

    Args:
        approx: Approximation coefficients.
        detail: Detail coefficients.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode the decomposition used.
        axis: Axis that was transformed.
        output_size: Length of the reconstructed axis.
    """
    sizes = None if output_size is None else (output_size,)
    return idwtn({"a": approx, "d": detail}, wavelet, mode, (axis,), sizes)

idwtn

idwtn(
    coeffs: dict[str, torch.Tensor],
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
    output_size: Sequence[int] | None = None,
) -> torch.Tensor

Invert a single-level separable wavelet transform.

Parameters:

Name	Type	Description	Default
coeffs	dict[str, torch.Tensor]	Subbands keyed by their 'a'/'d' names.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode the decomposition used.	'zero'
axes	Sequence[int] | None	Axes that were transformed; the trailing one if omitted.	None
output_size	Sequence[int] | None	Length of each transformed axis in the reconstruction; an even length is assumed if omitted.	None

Raises:

Type	Description
ValueError	If a subband is missing or the subbands disagree in shape.

Source code in src/chuchichaestli/dwt/functional.py

def idwtn(
    coeffs: dict[str, torch.Tensor],
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
    output_size: Sequence[int] | None = None,
) -> torch.Tensor:
    """Invert a single-level separable wavelet transform.

    Args:
        coeffs: Subbands keyed by their `'a'`/`'d'` names.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode the decomposition used.
        axes: Axes that were transformed; the trailing one if omitted.
        output_size: Length of each transformed axis in the reconstruction;
            an even length is assumed if omitted.

    Raises:
        ValueError: If a subband is missing or the subbands disagree in shape.
    """
    sample = next(iter(coeffs.values()))
    axes = _resolve_axes(sample.ndim, axes)
    keys = subband_keys(len(axes))
    missing = [key for key in keys if key not in coeffs]
    if missing:
        raise ValueError(
            f"Missing subband(s) {missing} for a {len(axes)}-dimensional transform;"
            f" expected {list(keys)}."
        )
    shapes = {tuple(coeffs[key].shape) for key in keys}
    if len(shapes) != 1:
        raise ValueError(f"All subbands must have the same shape; got {sorted(shapes)}.")

    wavelet = as_wavelet(wavelet)
    stacked = [_fold(as_inexact(coeffs[key]), axes) for key in keys]
    lead, perm = stacked[0][1], stacked[0][2]
    bands = [band for band, _, _ in stacked]
    h = _stack_bands(bands)
    _, _, rec_lo, rec_hi = wavelet.filters(h.dtype, h.device)

    sizes = _reconstruction_sizes(h.shape[2:], wavelet.filter_len, mode, output_size)
    from chuchichaestli.dwt import _ext

    if _ext.idwt_nd_applies(h.device, len(axes)) and _ext.kernels_available(
        h.device, h.dtype
    ):
        filter_len = wavelet.filter_len
        trim = filter_len // 2 - 1 if mode == "periodization" else filter_len - 2
        fused = _ext.idwt_nd(
            h, rec_lo, rec_hi, mode, (trim,) * len(axes), tuple(sizes)
        )
        return _unfold(fused[:, 0], lead, perm)
    # Decomposition appends one character per axis, so the axis transformed last
    # varies fastest and its band pairs are adjacent.
    for axis in reversed(range(len(axes))):
        h = _reconstruct(h, rec_lo, rec_hi, axis, mode, sizes[axis])
    return _unfold(h[:, 0], lead, perm)

subband_keys cached

subband_keys(dimensions: int) -> tuple[str, ...]

Subband names of a dimensions-dimensional transform, in output order.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of transformed axes.	required

Source code in src/chuchichaestli/dwt/functional.py

@lru_cache(maxsize=64)
def subband_keys(dimensions: int) -> tuple[str, ...]:
    """Subband names of a `dimensions`-dimensional transform, in output order.

    Args:
        dimensions: Number of transformed axes.
    """
    keys = ("",)
    for _ in range(dimensions):
        keys = tuple(k + branch for k in keys for branch in "ad")
    return keys

wavedec

wavedec(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    level: int | None = None,
    axis: int = -1,
) -> list[torch.Tensor]

Multi-level wavelet transform along one axis.

Parameters:

Name	Type	Description	Default
data	torch.Tensor	Input tensor of any rank.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode.	'zero'
level	int | None	Number of levels; the maximum the axis supports if omitted.	None
axis	int	Axis to transform.	-1

Returns:

Type	Description
list[torch.Tensor]	[approx, detail_level, ..., detail_1], coarsest first.

Source code in src/chuchichaestli/dwt/functional.py

def wavedec(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    level: int | None = None,
    axis: int = -1,
) -> list[torch.Tensor]:
    """Multi-level wavelet transform along one axis.

    Args:
        data: Input tensor of any rank.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode.
        level: Number of levels; the maximum the axis supports if omitted.
        axis: Axis to transform.

    Returns:
        `[approx, detail_level, ..., detail_1]`, coarsest first.
    """
    parts = wavedecn(data, wavelet, mode, level, (axis,))
    return [parts[0]] + [band["d"] for band in parts[1:]]

wavedecn

wavedecn(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    level: int | None = None,
    axes: Sequence[int] | None = None,
) -> list

Multi-level separable wavelet transform, recursing on the approximation.

Parameters:

Name	Type	Description	Default
data	torch.Tensor	Input tensor of any rank.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode.	'zero'
level	int | None	Number of levels; the maximum the shortest axis supports if omitted.	None
axes	Sequence[int] | None	Axes to transform; the trailing one if omitted.	None

Returns:

Type	Description
list	[approx, details_level, ..., details_1], coarsest first, each details
list	holding every subband but the all-approximation one.

Raises:

Type	Description
ValueError	If level is negative.

Source code in src/chuchichaestli/dwt/functional.py

def wavedecn(
    data: torch.Tensor,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    level: int | None = None,
    axes: Sequence[int] | None = None,
) -> list:
    """Multi-level separable wavelet transform, recursing on the approximation.

    Args:
        data: Input tensor of any rank.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode.
        level: Number of levels; the maximum the shortest axis supports if omitted.
        axes: Axes to transform; the trailing one if omitted.

    Returns:
        `[approx, details_level, ..., details_1]`, coarsest first, each `details`
        holding every subband but the all-approximation one.

    Raises:
        ValueError: If `level` is negative.
    """
    data = as_inexact(data)
    axes = _resolve_axes(data.ndim, axes)
    wavelet = as_wavelet(wavelet)
    if level is None:
        level = min(
            dwt_max_level(data.shape[axis], wavelet.dec_len) for axis in axes
        )
        level = max(level, 1)
    if level < 0:
        raise ValueError(f"The number of levels must not be negative; got {level}.")

    from chuchichaestli.dwt import _ext

    approx_key = "a" * len(axes)
    keys = subband_keys(len(axes))
    result: list = []

    folded, lead, perm = _fold(data, axes)
    shapes = [tuple(folded.shape[2:])]
    for _ in range(level - 1):
        shapes.append(
            tuple(dwt_coeff_len(n, wavelet.dec_len, mode) for n in shapes[-1])
        )
    compiled = level >= 1 and not (
        torch.is_grad_enabled() and data.requires_grad
    ) and _ext.kernels_available(folded.device, folded.dtype)

    # the kernel runs the recursion: the Haar form where every level suits it,
    # and the general one otherwise
    levels = None
    if compiled:
        if all(
            _ext.fused_haar_applies(wavelet, mode, shape) for shape in shapes
        ):
            levels = _ext.haar_wavedec(folded, len(axes), level)
        elif _ext.fused_recursion_applies(mode, shapes):
            dec_lo, dec_hi, _, _ = wavelet.filters(folded.dtype, folded.device)
            levels = _ext.wavedec_axes(folded, dec_lo, dec_hi, mode, level)
    if levels is not None:
        for stacked in levels:
            bands = {
                key: _unfold(stacked[:, i], lead, perm)
                for i, key in enumerate(keys)
            }
            approx = bands.pop(approx_key)
            result.append(bands)
        result.append(approx)
        return result[::-1]

    approx = data
    for _ in range(level):
        bands = dwtn(approx, wavelet, mode, axes)
        approx = bands.pop(approx_key)
        result.append(bands)
    result.append(approx)
    return result[::-1]

waverec

waverec(
    coeffs: Sequence[torch.Tensor],
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
    output_size: Sequence[int] | None = None,
) -> torch.Tensor

Invert a multi-level wavelet transform along one axis.

Parameters:

Name	Type	Description	Default
coeffs	Sequence[torch.Tensor]	[approx, detail_level, ..., detail_1], as wavedec returns.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode the decomposition used.	'zero'
axis	int	Axis that was transformed.	-1
output_size	Sequence[int] | None	Length of the axis at each level, coarsest level first.	None

Raises:

Type	Description
ValueError	If coeffs is empty.

Source code in src/chuchichaestli/dwt/functional.py

def waverec(
    coeffs: Sequence[torch.Tensor],
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axis: int = -1,
    output_size: Sequence[int] | None = None,
) -> torch.Tensor:
    """Invert a multi-level wavelet transform along one axis.

    Args:
        coeffs: `[approx, detail_level, ..., detail_1]`, as `wavedec` returns.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode the decomposition used.
        axis: Axis that was transformed.
        output_size: Length of the axis at each level, coarsest level first.

    Raises:
        ValueError: If `coeffs` is empty.
    """
    if not coeffs:
        raise ValueError("`coeffs` must hold at least the approximation band.")
    sizes = None if output_size is None else [(n,) for n in output_size]
    parts = [coeffs[0]] + [{"d": band} for band in coeffs[1:]]
    return waverecn(parts, wavelet, mode, (axis,), sizes)

waverecn

waverecn(
    coeffs: list,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
    output_size: Sequence[Sequence[int]] | None = None,
) -> torch.Tensor

Invert a multi-level separable wavelet transform.

Parameters:

Name	Type	Description	Default
coeffs	list	[approx, details_level, ..., details_1], as wavedecn returns.	required
wavelet	str | Wavelet	Wavelet, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode the decomposition used.	'zero'
axes	Sequence[int] | None	Axes that were transformed; the trailing one if omitted.	None
output_size	Sequence[Sequence[int]] | None	Per level, the length of each transformed axis, coarsest level first.	None

Raises:

Type	Description
ValueError	If coeffs is empty or malformed.

Source code in src/chuchichaestli/dwt/functional.py

def waverecn(
    coeffs: list,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "zero",
    axes: Sequence[int] | None = None,
    output_size: Sequence[Sequence[int]] | None = None,
) -> torch.Tensor:
    """Invert a multi-level separable wavelet transform.

    Args:
        coeffs: `[approx, details_level, ..., details_1]`, as `wavedecn` returns.
        wavelet: Wavelet, by name or as a `Wavelet`.
        mode: Signal extension mode the decomposition used.
        axes: Axes that were transformed; the trailing one if omitted.
        output_size: Per level, the length of each transformed axis, coarsest
            level first.

    Raises:
        ValueError: If `coeffs` is empty or malformed.
    """
    if not coeffs:
        raise ValueError("`coeffs` must hold at least the approximation band.")
    approx, details = coeffs[0], coeffs[1:]
    if not details:
        return approx
    axes = _resolve_axes(approx.ndim, axes)
    approx_key = "a" * len(axes)
    for i, band in enumerate(details):
        sizes = None if output_size is None else output_size[i]
        if sizes is None and band:
            approx = _trim_to_band(approx, next(iter(band.values())), axes)
        approx = idwtn({approx_key: approx, **band}, wavelet, mode, axes, sizes)
    return approx