Refactor of order methods for WaveletPacket2d

Make the functions creating the natural and frequency packet order
for WaveletPacket2d static methods of WaveletPacket2d.
This changes
* `get_natural_order` from instance to static function
* `get_frequency_order` from separate func to static function in
  the scope of WaveletPacket2d

Further, to make both methods consistent, `get_frequency_order`
now returns concatenated strings instead of a tuple
of single char strings.

Author: felixblanke

examples/deepfake_analysis/packet_plot.py

@@ -11,47 +11,6 @@
 import ptwt
 
 
-def get_freq_order(level: int):
-    """Get the frequency order for a given packet decomposition level.
-    Adapted from:
-    https://github.com/PyWavelets/pywt/blob/master/pywt/_wavelet_packets.py
-    The code elements denote the filter application order. The filters
-    are named following the pywt convention as:
-    a - LL, low-low coefficients
-    h - LH, low-high coefficients
-    v - HL, high-low coefficients
-    d - HH, high-high coefficients
-    """
-    wp_natural_path = list(product(["a", "h", "v", "d"], repeat=level))
-
-    def _get_graycode_order(level, x="a", y="d"):
-        graycode_order = [x, y]
-        for _ in range(level - 1):
-            graycode_order = [x + path for path in graycode_order] + [
-                y + path for path in graycode_order[::-1]
-            ]
-        return graycode_order
-
-    def _expand_2d_path(path):
-        expanded_paths = {"d": "hh", "h": "hl", "v": "lh", "a": "ll"}
-        return (
-            "".join([expanded_paths[p][0] for p in path]),
-            "".join([expanded_paths[p][1] for p in path]),
-        )
-
-    nodes: dict = {}
-    for (row_path, col_path), node in [
-        (_expand_2d_path(node), node) for node in wp_natural_path
-    ]:
-        nodes.setdefault(row_path, {})[col_path] = node
-    graycode_order = _get_graycode_order(level, x="l", y="h")
-    nodes_list: list = [nodes[path] for path in graycode_order if path in nodes]
-    wp_frequency_path = []
-    for row in nodes_list:
-        wp_frequency_path.append([row[path] for path in graycode_order if path in row])
-    return wp_frequency_path, wp_natural_path
-
-
 def generate_frequency_packet_image(packet_array: np.ndarray, degree: int):
     """Create a ready-to-polt image with frequency-order packages.
        Given a packet array in natural order, creat an image which is
@@ -63,7 +22,8 @@ def generate_frequency_packet_image(packet_array: np.ndarray, degree: int):
     Returns:
         [np.ndarray]: The image of shape [original_height, original_width]
     """
-    wp_freq_path, wp_natural_path = get_freq_order(degree)
+    wp_freq_path = ptwt.WaveletPacket2D.get_freq_order(degree)
+    wp_natural_path = ptwt.WaveletPacket2D.get_natural_order(degree)
 
     image = []
     # go through the rows.
@@ -107,7 +67,8 @@ def load_images(path: str) -> list:
 
 
 if __name__ == "__main__":
-    frequency_path, natural_path = get_freq_order(level=3)
+    freq_path = ptwt.WaveletPacket2D.get_freq_order(level=3)
+    frequency_path = ptwt.WaveletPacket2D.get_natural_order(level=3)
     print("Loading ffhq images:")
     ffhq_images = load_images("./ffhq_style_gan/source_data/A_ffhq")
     print("processing ffhq")

src/ptwt/packets.py

@@ -122,7 +122,7 @@ def __init__(
 
     def transform(
         self, data: torch.Tensor, maxlevel: Optional[int] = None
-    ) -> "WaveletPacket":
+    ) -> WaveletPacket:
         """Calculate the 1d wavelet packet transform for the input data.
 
         Args:
@@ -139,7 +139,7 @@ def transform(
         self._recursive_dwt(data, level=0, path="")
         return self
 
-    def reconstruct(self) -> "WaveletPacket":
+    def reconstruct(self) -> WaveletPacket:
         """Recursively reconstruct the input starting from the leaf nodes.
 
         Reconstruction replaces the input data originally assigned to this object.
@@ -326,7 +326,7 @@ def __init__(
 
     def transform(
         self, data: torch.Tensor, maxlevel: Optional[int] = None
-    ) -> "WaveletPacket2D":
+    ) -> WaveletPacket2D:
         """Calculate the 2d wavelet packet transform for the input data.
 
            The transform function allows reusing the same object.
@@ -350,7 +350,7 @@ def transform(
         self._recursive_dwt2d(data, level=0, path="")
         return self
 
-    def reconstruct(self) -> "WaveletPacket2D":
+    def reconstruct(self) -> WaveletPacket2D:
         """Recursively reconstruct the input starting from the leaf nodes.
 
         Note:
@@ -364,7 +364,7 @@ def reconstruct(self) -> "WaveletPacket2D":
             )
 
         for level in reversed(range(self.maxlevel)):
-            for node in self.get_natural_order(level):
+            for node in WaveletPacket2D.get_natural_order(level):
                 data_a = self[node + "a"]
                 data_h = self[node + "h"]
                 data_v = self[node + "v"]
@@ -386,17 +386,6 @@ def reconstruct(self) -> "WaveletPacket2D":
                 self[node] = rec
         return self
 
-    def get_natural_order(self, level: int) -> list[str]:
-        """Get the natural ordering for a given decomposition level.
-
-        Args:
-            level (int): The decomposition level.
-
-        Returns:
-            A list with the filter order strings.
-        """
-        return ["".join(p) for p in product(["a", "h", "v", "d"], repeat=level)]
-
     def _get_wavedec(self, shape: tuple[int, ...]) -> Callable[
         [torch.Tensor],
         WaveletCoeff2d,
@@ -525,54 +514,68 @@ def __getitem__(self, key: str) -> torch.Tensor:
             )
         return super().__getitem__(key)
 
+    @staticmethod
+    def get_natural_order(level: int) -> list[str]:
+        """Get the natural ordering for a given decomposition level.
 
-def get_freq_order(level: int) -> list[list[tuple[str, ...]]]:
-    """Get the frequency order for a given packet decomposition level.
+        Args:
+            level (int): The decomposition level.
 
-    Use this code to create two-dimensional frequency orderings.
+        Returns:
+            A list with the filter order strings.
+        """
+        return ["".join(p) for p in product(["a", "h", "v", "d"], repeat=level)]
 
-    Args:
-        level (int): The number of decomposition scales.
+    @staticmethod
+    def get_freq_order(level: int) -> list[list[str]]:
+        """Get the frequency order for a given packet decomposition level.
 
-    Returns:
-        A list with the tree nodes in frequency order.
-
-    Note:
-        Adapted from:
-        https://github.com/PyWavelets/pywt/blob/master/pywt/_wavelet_packets.py
-
-        The code elements denote the filter application order. The filters
-        are named following the pywt convention as:
-        a - LL, low-low coefficients
-        h - LH, low-high coefficients
-        v - HL, high-low coefficients
-        d - HH, high-high coefficients
-    """
-    wp_natural_path = product(["a", "h", "v", "d"], repeat=level)
+        Use this code to create two-dimensional frequency orderings.
 
-    def _get_graycode_order(level: int, x: str = "a", y: str = "d") -> list[str]:
-        graycode_order = [x, y]
-        for _ in range(level - 1):
-            graycode_order = [x + path for path in graycode_order] + [
-                y + path for path in graycode_order[::-1]
-            ]
-        return graycode_order
-
-    def _expand_2d_path(path: tuple[str, ...]) -> tuple[str, str]:
-        expanded_paths = {"d": "hh", "h": "hl", "v": "lh", "a": "ll"}
-        return (
-            "".join([expanded_paths[p][0] for p in path]),
-            "".join([expanded_paths[p][1] for p in path]),
-        )
-
-    nodes_dict: dict[str, dict[str, tuple[str, ...]]] = {}
-    for (row_path, col_path), node in [
-        (_expand_2d_path(node), node) for node in wp_natural_path
-    ]:
-        nodes_dict.setdefault(row_path, {})[col_path] = node
-    graycode_order = _get_graycode_order(level, x="l", y="h")
-    nodes = [nodes_dict[path] for path in graycode_order if path in nodes_dict]
-    result = []
-    for row in nodes:
-        result.append([row[path] for path in graycode_order if path in row])
-    return result
+        Args:
+            level (int): The number of decomposition scales.
+
+        Returns:
+            A list with the tree nodes in frequency order.
+
+        Note:
+            Adapted from:
+            https://github.com/PyWavelets/pywt/blob/master/pywt/_wavelet_packets.py
+
+            The code elements denote the filter application order. The filters
+            are named following the pywt convention as:
+            a - LL, low-low coefficients
+            h - LH, low-high coefficients
+            v - HL, high-low coefficients
+            d - HH, high-high coefficients
+        """
+        wp_natural_path = product(["a", "h", "v", "d"], repeat=level)
+
+        def _get_graycode_order(level: int, x: str = "a", y: str = "d") -> list[str]:
+            graycode_order = [x, y]
+            for _ in range(level - 1):
+                graycode_order = [x + path for path in graycode_order] + [
+                    y + path for path in graycode_order[::-1]
+                ]
+            return graycode_order
+
+        def _expand_2d_path(path: tuple[str, ...]) -> tuple[str, str]:
+            expanded_paths = {"d": "hh", "h": "hl", "v": "lh", "a": "ll"}
+            return (
+                "".join([expanded_paths[p][0] for p in path]),
+                "".join([expanded_paths[p][1] for p in path]),
+            )
+
+        nodes_dict: dict[str, dict[str, tuple[str, ...]]] = {}
+        for (row_path, col_path), node in [
+            (_expand_2d_path(node), node) for node in wp_natural_path
+        ]:
+            nodes_dict.setdefault(row_path, {})[col_path] = node
+        graycode_order = _get_graycode_order(level, x="l", y="h")
+        nodes = [nodes_dict[path] for path in graycode_order if path in nodes_dict]
+        result = []
+        for row in nodes:
+            result.append(
+                ["".join(row[path]) for path in graycode_order if path in row]
+            )
+        return result

tests/test_packets.py

@@ -12,7 +12,7 @@
 from scipy import datasets
 
 from ptwt.constants import ExtendedBoundaryMode
-from ptwt.packets import WaveletPacket, WaveletPacket2D, get_freq_order
+from ptwt.packets import WaveletPacket, WaveletPacket2D
 
 
 def _compare_trees1(
@@ -236,7 +236,7 @@ def test_freq_order(level: int, wavelet_str: str, pywt_boundary: str) -> None:
     )
     # Get the full decomposition
     freq_tree = wp_tree.get_level(level, "freq")
-    freq_order = get_freq_order(level)
+    freq_order = WaveletPacket2D.get_freq_order(level)
 
     for order_list, tree_list in zip(freq_tree, freq_order):
         for order_el, tree_el in zip(order_list, tree_list):