Add 1d return type

Author: felixblanke

docs/ref/return-types.rst

@@ -5,6 +5,12 @@
 Wavelet transform return types
 ==============================
 
+Transforms in one dimension
+---------------------------
+
+.. autodata:: WaveletCoeff1d
+
+
 Transforms in two dimensions
 ----------------------------
 

src/ptwt/constants.py

@@ -1,5 +1,6 @@
 """Constants and types used throughout the PyTorch Wavelet Toolbox."""
 
+from collections.abc import Sequence
 from typing import Literal, NamedTuple, Union
 
 import torch
@@ -62,6 +63,27 @@
 """
 
 
+# Note: This data structure was chosen to follow pywt's conventions
+WaveletCoeff1d: TypeAlias = Sequence[torch.Tensor]
+"""Type alias for 1d wavelet transform results.
+
+This type alias represents the result of a 1d wavelet transform
+with :math:`n` levels as a sequence::
+
+    [cA_n, cD_n, cD_n-1, …, cD1]
+
+of :math:`n + 1` tensors.
+The first entry of the sequence (``cA_n``) is the approximation coefficient tensor.
+The following entries (``cD_n`` - ``cD1``) are the detail coefficient tensors
+of the respective level.
+
+Note that this type always contains an approximation coefficient tensor but does not
+necesseraily contain any detail coefficients.
+
+Alias of ``Sequence[torch.Tensor]``
+"""
+
+
 class WaveletDetailTuple2d(NamedTuple):
     """Detail coefficients of a 2d wavelet transform for a given level.
 

src/ptwt/conv_transform.py

@@ -21,7 +21,7 @@
     _pad_symmetric,
     _unfold_axes,
 )
-from .constants import BoundaryMode, WaveletCoeff2d
+from .constants import BoundaryMode, WaveletCoeff1d, WaveletCoeff2d
 
 
 def _create_tensor(
@@ -361,12 +361,13 @@ def wavedec(
 
 
 def waverec(
-    coeffs: Sequence[torch.Tensor], wavelet: Union[Wavelet, str], axis: int = -1
+    coeffs: WaveletCoeff1d, wavelet: Union[Wavelet, str], axis: int = -1
 ) -> torch.Tensor:
     """Reconstruct a 1d signal from wavelet coefficients.
 
     Args:
-        coeffs (Sequence): The wavelet coefficient sequence produced by wavedec.
+        coeffs: The wavelet coefficient sequence produced by the forward transform
+            :func:`wavedec`. See :data:`ptwt.constants.WaveletCoeff1d`.
         wavelet (Wavelet or str): A pywt wavelet compatible object or
             the name of a pywt wavelet.
             Refer to the output from ``pywt.wavelist(kind='discrete')``

src/ptwt/matmul_transform.py

@@ -8,7 +8,6 @@
 """
 
 import sys
-from collections.abc import Sequence
 from typing import Optional, Union
 
 import numpy as np
@@ -21,7 +20,7 @@
     _is_dtype_supported,
     _unfold_axes,
 )
-from .constants import OrthogonalizeMethod, PaddingMode
+from .constants import OrthogonalizeMethod, PaddingMode, WaveletCoeff1d
 from .conv_transform import (
     _get_filter_tensors,
     _postprocess_result_list_dec1d,
@@ -595,12 +594,12 @@ def _construct_synthesis_matrices(
             self.ifwt_matrix_list.append(sn)
             curr_length = curr_length // 2
 
-    def __call__(self, coefficients: Sequence[torch.Tensor]) -> torch.Tensor:
+    def __call__(self, coefficients: WaveletCoeff1d) -> torch.Tensor:
         """Run the synthesis or inverse matrix fwt.
 
         Args:
-            coefficients (Sequence[torch.Tensor]): The coefficients produced
-                by the forward transform.
+            coefficients: The coefficients produced by the forward transform
+                :data:`MatrixWavedec`. See :data:`ptwt.constants.WaveletCoeff1d`.
 
         Returns:
             The input signal reconstruction.

src/ptwt/packets.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import collections
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable
 from functools import partial
 from itertools import product
 from typing import TYPE_CHECKING, Literal, Optional, Union, overload
@@ -17,6 +17,7 @@
     ExtendedBoundaryMode,
     OrthogonalizeMethod,
     PacketNodeOrder,
+    WaveletCoeff1d,
     WaveletCoeff2d,
     WaveletCoeffNd,
     WaveletDetailTuple2d,
@@ -230,7 +231,7 @@ def _get_wavedec(
     def _get_waverec(
         self,
         length: int,
-    ) -> Callable[[Sequence[torch.Tensor]], torch.Tensor]:
+    ) -> Callable[[WaveletCoeff1d], torch.Tensor]:
         if self.mode == "boundary":
             if length not in self._matrix_waverec_dict.keys():
                 self._matrix_waverec_dict[length] = MatrixWaverec(

src/ptwt/stationary_transform.py

@@ -8,6 +8,7 @@
 import torch.nn.functional as F  # noqa:N812
 
 from ._util import Wavelet, _as_wavelet, _unfold_axes
+from .constants import WaveletCoeff1d
 from .conv_transform import (
     _get_filter_tensors,
     _postprocess_result_list_dec1d,
@@ -109,15 +110,15 @@ def swt(
 
 
 def iswt(
-    coeffs: Sequence[torch.Tensor],
+    coeffs: WaveletCoeff1d,
     wavelet: Union[pywt.Wavelet, str],
     axis: Optional[int] = -1,
 ) -> torch.Tensor:
     """Invert a 1d stationary wavelet transform.
 
     Args:
-        coeffs (Sequence[torch.Tensor]): The coefficients as computed
-            by the swt function.
+        coeffs: The wavelet coefficient sequence produced by the forward transform
+            :func:`swt`. See :data:`ptwt.constants.WaveletCoeff1d`.
         wavelet (Wavelet or str): A pywt wavelet compatible object or
             the name of a pywt wavelet, as used in the forward transform.
         axis (int, optional): The axis the forward trasform was computed over.