chuchichaestli.dwt.modes

Signal extension modes for the boundary handling of the wavelet transform.

Functions:

Name	Description
extension_indices	Build the gather that extends a signal of length n past both ends.
pad_signal	Extend a tensor along one axis by the given signal extension mode.
pad_signal_adjoint	Transpose of pad_signal, scattering a gradient back onto the samples.

extension_indices cached

extension_indices(
    n: int,
    pad_lo: int,
    pad_hi: int,
    mode: ExtensionModeTypes,
) -> tuple[
    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]

Build the gather that extends a signal of length n past both ends.

The result is cached on CPU; callers move it to the device they need.

Parameters:

Name	Type	Description	Default
n	int	Length of the signal.	required
pad_lo	int	Number of samples prepended.	required
pad_hi	int	Number of samples appended.	required
mode	ExtensionModeTypes	Signal extension mode.	required

Raises:

Type	Description
ValueError	If n is not positive or a padding is negative.

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

@lru_cache(maxsize=512)
def extension_indices(
    n: int, pad_lo: int, pad_hi: int, mode: ExtensionModeTypes
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Build the gather that extends a signal of length `n` past both ends.

    The result is cached on CPU; callers move it to the device they need.

    Args:
        n: Length of the signal.
        pad_lo: Number of samples prepended.
        pad_hi: Number of samples appended.
        mode: Signal extension mode.

    Raises:
        ValueError: If `n` is not positive or a padding is negative.
    """
    if n < 1:
        raise ValueError(f"A signal needs at least one sample; got {n}.")
    if pad_lo < 0 or pad_hi < 0:
        raise ValueError(
            f"Signal extension widths must not be negative; got ({pad_lo}, {pad_hi})."
        )
    entries = [_sample(i, n, mode) for i in range(-pad_lo, n + pad_hi)]
    sign, index, lo, hi = zip(*entries)
    return (
        torch.tensor(sign, dtype=torch.float64),
        torch.tensor(index, dtype=torch.int64),
        torch.tensor(lo, dtype=torch.float64),
        torch.tensor(hi, dtype=torch.float64),
    )

pad_signal

pad_signal(
    x: torch.Tensor,
    axis: int,
    pad_lo: int,
    pad_hi: int,
    mode: ExtensionModeTypes = "zero",
) -> torch.Tensor

Extend a tensor along one axis by the given signal extension mode.

Parameters:

Name	Type	Description	Default
x	torch.Tensor	Input tensor.	required
axis	int	Axis to extend, given as a non-negative or negative index.	required
pad_lo	int	Number of samples prepended.	required
pad_hi	int	Number of samples appended.	required
mode	ExtensionModeTypes	Signal extension mode.	'zero'

Raises:

Type	Description
ValueError	If a padding is negative.

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

def pad_signal(
    x: torch.Tensor,
    axis: int,
    pad_lo: int,
    pad_hi: int,
    mode: ExtensionModeTypes = "zero",
) -> torch.Tensor:
    """Extend a tensor along one axis by the given signal extension mode.

    Args:
        x: Input tensor.
        axis: Axis to extend, given as a non-negative or negative index.
        pad_lo: Number of samples prepended.
        pad_hi: Number of samples appended.
        mode: Signal extension mode.

    Raises:
        ValueError: If a padding is negative.
    """
    if pad_lo < 0 or pad_hi < 0:
        raise ValueError(
            f"Signal extension widths must not be negative; got ({pad_lo}, {pad_hi})."
        )
    if pad_lo == 0 and pad_hi == 0:
        return x
    axis = axis % x.ndim
    if mode == "zero":
        # `F.pad` reads its widths from the last axis backwards.
        widths = [0, 0] * (x.ndim - 1 - axis) + [pad_lo, pad_hi]
        return F.pad(x, widths, mode="constant", value=0.0)

    n = x.shape[axis]
    sign, index, lo, hi = extension_indices(n, pad_lo, pad_hi, mode)
    index = index.to(x.device)
    out = x.index_select(axis, index)
    if not bool(torch.all(sign == 1.0)):
        out = out * view_along_axis(sign.to(x.dtype), x.ndim, axis).to(x.device)
    if bool(torch.any(lo != 0.0)) or bool(torch.any(hi != 0.0)):
        edge_lo = x.narrow(axis, 0, 1)
        edge_hi = x.narrow(axis, n - 1, 1)
        out = out + edge_lo * view_along_axis(lo.to(x.dtype), x.ndim, axis).to(x.device)
        out = out + edge_hi * view_along_axis(hi.to(x.dtype), x.ndim, axis).to(x.device)
    return out

pad_signal_adjoint

pad_signal_adjoint(
    grad: torch.Tensor,
    axis: int,
    pad_lo: int,
    pad_hi: int,
    mode: ExtensionModeTypes = "zero",
    length: int | None = None,
) -> torch.Tensor

Transpose of pad_signal, scattering a gradient back onto the samples.

The extension gathers, so its transpose scatters: every extended sample returns its share to the source it was read from, and to the two edges the anchored modes lean on.

Parameters:

Name	Type	Description	Default
grad	torch.Tensor	Gradient with respect to the extended tensor.	required
axis	int	Axis that was extended.	required
pad_lo	int	Number of samples that were prepended.	required
pad_hi	int	Number of samples that were appended.	required
mode	ExtensionModeTypes	Signal extension mode that was used.	'zero'
length	int | None	Length of the axis before extension; derived if omitted.	None

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

def pad_signal_adjoint(
    grad: torch.Tensor,
    axis: int,
    pad_lo: int,
    pad_hi: int,
    mode: ExtensionModeTypes = "zero",
    length: int | None = None,
) -> torch.Tensor:
    """Transpose of `pad_signal`, scattering a gradient back onto the samples.

    The extension gathers, so its transpose scatters: every extended sample
    returns its share to the source it was read from, and to the two edges the
    anchored modes lean on.

    Args:
        grad: Gradient with respect to the extended tensor.
        axis: Axis that was extended.
        pad_lo: Number of samples that were prepended.
        pad_hi: Number of samples that were appended.
        mode: Signal extension mode that was used.
        length: Length of the axis before extension; derived if omitted.
    """
    axis = axis % grad.ndim
    if pad_lo == 0 and pad_hi == 0:
        return grad
    n = grad.shape[axis] - pad_lo - pad_hi if length is None else length
    if mode == "zero":
        return grad.narrow(axis, pad_lo, n)

    sign, index, lo, hi = extension_indices(n, pad_lo, pad_hi, mode)
    index = index.to(grad.device)
    contribution = grad * view_along_axis(sign.to(grad.dtype), grad.ndim, axis).to(
        grad.device
    )
    shape = list(grad.shape)
    shape[axis] = n
    out = torch.zeros(shape, dtype=grad.dtype, device=grad.device)
    out = out.index_add(axis, index, contribution)
    if bool(torch.any(lo != 0.0)) or bool(torch.any(hi != 0.0)):
        weights_lo = view_along_axis(lo.to(grad.dtype), grad.ndim, axis).to(grad.device)
        weights_hi = view_along_axis(hi.to(grad.dtype), grad.ndim, axis).to(grad.device)
        edge_lo = (grad * weights_lo).sum(dim=axis, keepdim=True)
        edge_hi = (grad * weights_hi).sum(dim=axis, keepdim=True)
        # a single-sample axis makes both edges the same one, which `index_add`
        # handles by accumulating the two contributions
        edges = torch.tensor([0, n - 1], device=grad.device)
        out = out.index_add(axis, edges, torch.cat((edge_lo, edge_hi), dim=axis))
    return out