chuchichaestli.data.transforms.wavelet

Discrete wavelet transforms over the trailing spatial dimensions.

Classes:

Name	Description
InvWaveletDecompose	Inverse of WaveletDecompose: reconstruct forward, decompose on revert.
InvWaveletPacket	Inverse of WaveletPacket: reconstruct forward, decompose on revert.
WaveletDecompose	Replace the trailing ndim axes by the subbands of a wavelet transform.
WaveletPacket	WaveletDecompose that keeps splitting every subband, not only the coarsest.

InvWaveletDecompose

InvWaveletDecompose(
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
)

Bases: WaveletDecompose

Inverse of WaveletDecompose: reconstruct forward, decompose on revert.

Methods:

Name	Description
get_inverse	Return the transform that decomposes what this one reconstructs.
make_params	Report the recorded spatial shape; the forward pass reads, not writes.
revert	Decompose the trailing spatial axes.
transform	Reconstruct the trailing spatial axes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def __init__(
    self,
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
):
    """Constructor.

    Args:
        wavelet: Wavelet to decompose with, by name or as a `Wavelet`.
        ndim: Number of trailing spatial dims to transform.
        mode: Signal extension mode; `'periodization'` is the only one that
            halves each axis exactly, and so the only one that does not grow
            the data.
        stack_dim: Position of the subband axis in the result; immediately
            before the spatial axes if omitted.
        sizes: Spatial shape the coefficients revert to, if it is known.

    Raises:
        ValueError: If `ndim` is out of range.
    """
    super().__init__()
    if not 1 <= ndim <= 3:
        raise ValueError(f"The transform runs over one to three axes; got {ndim}.")
    levels = 1
    self.wavelet = wavelet
    self.ndim = ndim
    self.mode = mode
    self.levels = levels
    self.stack_dim = -ndim - 1 if stack_dim is None else stack_dim
    self.sizes = None if sizes is None else tuple(int(n) for n in sizes)
    self._declared = sizes is not None
    self._coerced: Wavelet | None = None

get_inverse

get_inverse() -> WaveletDecompose

Return the transform that decomposes what this one reconstructs.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def get_inverse(self) -> "WaveletDecompose":
    """Return the transform that decomposes what this one reconstructs."""
    return WaveletDecompose(
        self.wavelet, self.ndim, self.mode, self.stack_dim, self.sizes
    )

make_params

make_params(flat_inputs: list[Any]) -> dict[str, Any]

Report the recorded spatial shape; the forward pass reads, not writes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
    """Report the recorded spatial shape; the forward pass reads, not writes."""
    return {"sizes": self.sizes}

revert

revert(x: Any, params: dict[str, Any] | None = None) -> Any

Decompose the trailing spatial axes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def revert(self, x: Any, params: dict[str, Any] | None = None) -> Any:
    """Decompose the trailing spatial axes."""
    self._record(tuple(int(n) for n in x.shape[-self.ndim :]))
    return self._decompose(x)

transform

transform(
    x: Any, params: dict[str, Any] | None = None
) -> Any

Reconstruct the trailing spatial axes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def transform(self, x: Any, params: dict[str, Any] | None = None) -> Any:
    """Reconstruct the trailing spatial axes."""
    return self._reconstruct(x)

InvWaveletPacket

InvWaveletPacket(
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    levels: int = 2,
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
)

Bases: WaveletPacket, InvWaveletDecompose

Inverse of WaveletPacket: reconstruct forward, decompose on revert.

Takes the level count from WaveletPacket and the swapped direction from InvWaveletDecompose, in that order.

Methods:

Name	Description
get_inverse	Return the transform that decomposes what this one reconstructs.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def __init__(
    self,
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    levels: int = 2,
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
):
    """Constructor.

    Args:
        wavelet: Wavelet to decompose with, by name or as a `Wavelet`.
        ndim: Number of trailing spatial dims to transform.
        mode: Signal extension mode.
        levels: Number of levels to decompose.
        stack_dim: Position of the subband axis in the result.
        sizes: Spatial shape the coefficients revert to, if it is known.

    Raises:
        ValueError: If `levels` is not positive.
    """
    super().__init__(wavelet, ndim, mode, stack_dim, sizes)
    if levels < 1:
        raise ValueError(f"A decomposition needs at least one level; got {levels}.")
    self.levels = levels

get_inverse

get_inverse() -> WaveletPacket

Return the transform that decomposes what this one reconstructs.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def get_inverse(self) -> "WaveletPacket":
    """Return the transform that decomposes what this one reconstructs."""
    return WaveletPacket(
        self.wavelet,
        self.ndim,
        self.mode,
        self.levels,
        self.stack_dim,
        self.sizes,
    )

WaveletDecompose

WaveletDecompose(
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
)

Bases: Transform

Replace the trailing ndim axes by the subbands of a wavelet transform.

The subbands are stacked on one new axis, so the result stays a single tensor and collates like any other field. Leading axes pass through. Only one level is taken: recursing on the coarsest subband alone would leave the levels at different resolutions, which no single tensor can hold. See WaveletPacket for the deeper decomposition that does keep one resolution.

Reverting needs the spatial shape the coefficients came from, which the transform records as it runs; pass sizes to fix it up front instead.

Constructor.

Parameters:

Name	Type	Description	Default
wavelet	str | Wavelet	Wavelet to decompose with, by name or as a Wavelet.	'haar'
ndim	int	Number of trailing spatial dims to transform.	2
mode	ExtensionModeTypes	Signal extension mode; 'periodization' is the only one that halves each axis exactly, and so the only one that does not grow the data.	'zero'
stack_dim	int | None	Position of the subband axis in the result; immediately before the spatial axes if omitted.	None
sizes	Sequence[int] | None	Spatial shape the coefficients revert to, if it is known.	None

Raises:

Type	Description
ValueError	If ndim is out of range.

Methods:

Name	Description
get_inverse	Return the transform that reconstructs what this one decomposes.
make_params	Record the spatial shape of the first tensor, shared by every leaf.
revert	Reconstruct the trailing spatial axes.
transform	Decompose the trailing spatial axes.

Attributes:

Name	Type	Description
axes	tuple[int, ...]	Axes the transform runs over, as trailing indices.
subbands	int	Number of subbands the transform produces.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def __init__(
    self,
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
):
    """Constructor.

    Args:
        wavelet: Wavelet to decompose with, by name or as a `Wavelet`.
        ndim: Number of trailing spatial dims to transform.
        mode: Signal extension mode; `'periodization'` is the only one that
            halves each axis exactly, and so the only one that does not grow
            the data.
        stack_dim: Position of the subband axis in the result; immediately
            before the spatial axes if omitted.
        sizes: Spatial shape the coefficients revert to, if it is known.

    Raises:
        ValueError: If `ndim` is out of range.
    """
    super().__init__()
    if not 1 <= ndim <= 3:
        raise ValueError(f"The transform runs over one to three axes; got {ndim}.")
    levels = 1
    self.wavelet = wavelet
    self.ndim = ndim
    self.mode = mode
    self.levels = levels
    self.stack_dim = -ndim - 1 if stack_dim is None else stack_dim
    self.sizes = None if sizes is None else tuple(int(n) for n in sizes)
    self._declared = sizes is not None
    self._coerced: Wavelet | None = None

axes property

axes: tuple[int, ...]

Axes the transform runs over, as trailing indices.

subbands property

subbands: int

Number of subbands the transform produces.

get_inverse

get_inverse() -> InvWaveletDecompose

Return the transform that reconstructs what this one decomposes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def get_inverse(self) -> "InvWaveletDecompose":
    """Return the transform that reconstructs what this one decomposes."""
    return InvWaveletDecompose(
        self.wavelet, self.ndim, self.mode, self.stack_dim, self.sizes
    )

make_params

make_params(flat_inputs: list[Any]) -> dict[str, Any]

Record the spatial shape of the first tensor, shared by every leaf.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
    """Record the spatial shape of the first tensor, shared by every leaf."""
    x = next(i for i in flat_inputs if isinstance(i, torch.Tensor))
    shape = tuple(int(n) for n in x.shape[-self.ndim :])
    self._record(shape)
    return {"sizes": shape}

revert

revert(x: Any, params: dict[str, Any] | None = None) -> Any

Reconstruct the trailing spatial axes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def revert(self, x: Any, params: dict[str, Any] | None = None) -> Any:
    """Reconstruct the trailing spatial axes."""
    return self._reconstruct(x)

transform

transform(
    x: Any, params: dict[str, Any] | None = None
) -> Any

Decompose the trailing spatial axes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def transform(self, x: Any, params: dict[str, Any] | None = None) -> Any:
    """Decompose the trailing spatial axes."""
    self._record(tuple(int(n) for n in x.shape[-self.ndim :]))
    return self._decompose(x)

WaveletPacket

WaveletPacket(
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    levels: int = 2,
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
)

Bases: WaveletDecompose

WaveletDecompose that keeps splitting every subband, not only the coarsest.

A packet decomposition of l levels leaves 2 ** (ndim * l) subbands, all at the same resolution, which is what makes the result a single tensor.

Constructor.

Parameters:

Name	Type	Description	Default
wavelet	str | Wavelet	Wavelet to decompose with, by name or as a Wavelet.	'haar'
ndim	int	Number of trailing spatial dims to transform.	2
mode	ExtensionModeTypes	Signal extension mode.	'zero'
levels	int	Number of levels to decompose.	2
stack_dim	int | None	Position of the subband axis in the result.	None
sizes	Sequence[int] | None	Spatial shape the coefficients revert to, if it is known.	None

Raises:

Type	Description
ValueError	If levels is not positive.

Methods:

Name	Description
get_inverse	Return the transform that reconstructs what this one decomposes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def __init__(
    self,
    wavelet: str | Wavelet = "haar",
    ndim: int = 2,
    mode: ExtensionModeTypes = "zero",
    levels: int = 2,
    stack_dim: int | None = None,
    sizes: Sequence[int] | None = None,
):
    """Constructor.

    Args:
        wavelet: Wavelet to decompose with, by name or as a `Wavelet`.
        ndim: Number of trailing spatial dims to transform.
        mode: Signal extension mode.
        levels: Number of levels to decompose.
        stack_dim: Position of the subband axis in the result.
        sizes: Spatial shape the coefficients revert to, if it is known.

    Raises:
        ValueError: If `levels` is not positive.
    """
    super().__init__(wavelet, ndim, mode, stack_dim, sizes)
    if levels < 1:
        raise ValueError(f"A decomposition needs at least one level; got {levels}.")
    self.levels = levels

get_inverse

get_inverse() -> InvWaveletPacket

Return the transform that reconstructs what this one decomposes.

Source code in src/chuchichaestli/data/transforms/wavelet.py

def get_inverse(self) -> "InvWaveletPacket":
    """Return the transform that reconstructs what this one decomposes."""
    return InvWaveletPacket(
        self.wavelet,
        self.ndim,
        self.mode,
        self.levels,
        self.stack_dim,
        self.sizes,
    )