Improve typing in JIT code

Author: felixblanke

examples/speed_tests/timeitconv_1d.py

@@ -9,26 +9,8 @@
 import ptwt
 
 
-class WaveletTuple(NamedTuple):
-    """Replaces namedtuple("Wavelet", ("dec_lo", "dec_hi", "rec_lo", "rec_hi"))."""
-
-    dec_lo: torch.Tensor
-    dec_hi: torch.Tensor
-    rec_lo: torch.Tensor
-    rec_hi: torch.Tensor
-
-
-def _set_up_wavelet_tuple(wavelet, dtype):
-    return WaveletTuple(
-        torch.tensor(wavelet.dec_lo).type(dtype),
-        torch.tensor(wavelet.dec_hi).type(dtype),
-        torch.tensor(wavelet.rec_lo).type(dtype),
-        torch.tensor(wavelet.rec_hi).type(dtype),
-    )
-
-
 def _jit_wavedec_fun(data, wavelet):
-    return ptwt.wavedec(data, wavelet, "periodic", level=10)
+    return ptwt.wavedec(data, wavelet, mode="periodic", level=10)
 
 
 if __name__ == "__main__":
@@ -56,7 +38,7 @@ def _jit_wavedec_fun(data, wavelet):
         end = time.perf_counter()
         ptwt_time_cpu.append(end - start)
 
-    wavelet = _set_up_wavelet_tuple(pywt.Wavelet("db5"), torch.float32)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(pywt.Wavelet("db5"), torch.float32)
     jit_wavedec = torch.jit.trace(
         _jit_wavedec_fun,
         (data, wavelet),
@@ -81,7 +63,7 @@ def _jit_wavedec_fun(data, wavelet):
         end = time.perf_counter()
         ptwt_time_gpu.append(end - start)
 
-    wavelet = _set_up_wavelet_tuple(pywt.Wavelet("db5"), torch.float32)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(pywt.Wavelet("db5"), torch.float32)
     jit_wavedec = torch.jit.trace(
         _jit_wavedec_fun,
         (data.cuda(), wavelet),

examples/speed_tests/timeitconv_2d.py

@@ -9,27 +9,9 @@
 import ptwt
 
 
-class WaveletTuple(NamedTuple):
-    """Replaces namedtuple("Wavelet", ("dec_lo", "dec_hi", "rec_lo", "rec_hi"))."""
-
-    dec_lo: torch.Tensor
-    dec_hi: torch.Tensor
-    rec_lo: torch.Tensor
-    rec_hi: torch.Tensor
-
-
-def _set_up_wavelet_tuple(wavelet, dtype):
-    return WaveletTuple(
-        torch.tensor(wavelet.dec_lo).type(dtype),
-        torch.tensor(wavelet.dec_hi).type(dtype),
-        torch.tensor(wavelet.rec_lo).type(dtype),
-        torch.tensor(wavelet.rec_hi).type(dtype),
-    )
-
-
-def _to_jit_wavedec_2(data, wavelet):
+def _to_jit_wavedec_2(data: torch.Tensor, wavelet) -> list[torch.Tensor]:
     """Ensure uniform datatypes in lists for the tracer.
-    Going from List[Union[torch.Tensor, List[torch.Tensor]]] to List[torch.Tensor]
+    Going from list[Union[torch.Tensor, list[torch.Tensor]]] to list[torch.Tensor]
     means we have to stack the lists in the output.
     """
     assert data.shape == (32, 1e3, 1e3), "Changing the chape requires re-tracing."
@@ -79,7 +61,7 @@ def _to_jit_wavedec_2(data, wavelet):
 
         ptwt_time_gpu.append(end - start)
 
-    wavelet = _set_up_wavelet_tuple(pywt.Wavelet("db5"), torch.float32)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(pywt.Wavelet("db5"), torch.float32)
     jit_wavedec = torch.jit.trace(
         _to_jit_wavedec_2,
         (data.cuda(), wavelet),

examples/speed_tests/timeitconv_2d_separable.py

@@ -10,31 +10,13 @@
 import ptwt
 
 
-class WaveletTuple(NamedTuple):
-    """Replaces namedtuple("Wavelet", ("dec_lo", "dec_hi", "rec_lo", "rec_hi"))."""
-
-    dec_lo: torch.Tensor
-    dec_hi: torch.Tensor
-    rec_lo: torch.Tensor
-    rec_hi: torch.Tensor
-
-
-def _set_up_wavelet_tuple(wavelet, dtype):
-    return WaveletTuple(
-        torch.tensor(wavelet.dec_lo).type(dtype),
-        torch.tensor(wavelet.dec_hi).type(dtype),
-        torch.tensor(wavelet.rec_lo).type(dtype),
-        torch.tensor(wavelet.rec_hi).type(dtype),
-    )
-
-
 def _to_jit_wavedec_2(data, wavelet):
     """Ensure uniform datatypes in lists for the tracer.
     Going from List[Union[torch.Tensor, List[torch.Tensor]]] to List[torch.Tensor]
     means we have to stack the lists in the output.
     """
     assert data.shape == (32, 1e3, 1e3), "Changing the chape requires re-tracing."
-    coeff = ptwt.fswavedec2(data, wavelet, "reflect", level=5)
+    coeff = ptwt.fswavedec2(data, wavelet, mode="reflect", level=5)
     coeff2 = []
     for c in coeff:
         if isinstance(c, torch.Tensor):
@@ -103,7 +85,7 @@ def _to_jit_wavedec_2(data, wavelet):
         end = time.perf_counter()
         ptwt_time_gpu.append(end - start)
 
-    wavelet = _set_up_wavelet_tuple(pywt.Wavelet("db5"), torch.float32)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(pywt.Wavelet("db5"), torch.float32)
     jit_wavedec = torch.jit.trace(
         _to_jit_wavedec_2,
         (data.cuda(), wavelet),

examples/speed_tests/timeitconv_3d.py

@@ -9,24 +9,6 @@
 import ptwt
 
 
-class WaveletTuple(NamedTuple):
-    """Replaces namedtuple("Wavelet", ("dec_lo", "dec_hi", "rec_lo", "rec_hi"))."""
-
-    dec_lo: torch.Tensor
-    dec_hi: torch.Tensor
-    rec_lo: torch.Tensor
-    rec_hi: torch.Tensor
-
-
-def _set_up_wavelet_tuple(wavelet, dtype):
-    return WaveletTuple(
-        torch.tensor(wavelet.dec_lo).type(dtype),
-        torch.tensor(wavelet.dec_hi).type(dtype),
-        torch.tensor(wavelet.rec_lo).type(dtype),
-        torch.tensor(wavelet.rec_hi).type(dtype),
-    )
-
-
 def _to_jit_wavedec_3(data, wavelet):
     """Ensure uniform datatypes in lists for the tracer.
 
@@ -85,7 +67,7 @@ def _to_jit_wavedec_3(data, wavelet):
         end = time.perf_counter()
         ptwt_time_gpu.append(end - start)
 
-    wavelet = _set_up_wavelet_tuple(pywt.Wavelet("db5"), torch.float32)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(pywt.Wavelet("db5"), torch.float32)
     jit_wavedec = torch.jit.trace(
         _to_jit_wavedec_3,
         (data.cuda(), wavelet),

src/ptwt/__init__.py

@@ -1,7 +1,7 @@
 """Differentiable and gpu enabled fast wavelet transforms in PyTorch."""
 
-from ._util import Wavelet
-from .constants import WaveletCoeff2d, WaveletCoeffNd, WaveletCoeff2dSeparable
+from ._util import Wavelet, WaveletTensorTuple
+from .constants import WaveletCoeff2d, WaveletCoeff2dSeparable, WaveletCoeffNd
 from .continuous_transform import cwt
 from .conv_transform import wavedec, waverec
 from .conv_transform_2 import wavedec2, waverec2

src/ptwt/_util.py

@@ -4,7 +4,7 @@
 
 import typing
 from collections.abc import Sequence
-from typing import Any, Callable, Optional, Protocol, Union, cast, overload
+from typing import Any, Callable, NamedTuple, Optional, Protocol, Union, cast, overload
 
 import numpy as np
 import pywt
@@ -38,6 +38,42 @@ def __len__(self) -> int:
         return len(self.dec_lo)
 
 
+class WaveletTensorTuple(NamedTuple):
+    """Named tuple containing the wavelet filter bank to use in JIT code."""
+
+    dec_lo: torch.Tensor
+    dec_hi: torch.Tensor
+    rec_lo: torch.Tensor
+    rec_hi: torch.Tensor
+
+    @property
+    def dec_len(self) -> int:
+        """Length of decomposition filters."""
+        return len(self.dec_lo)
+
+    @property
+    def rec_len(self) -> int:
+        """Length of reconstruction filters."""
+        return len(self.rec_lo)
+
+    @property
+    def filter_bank(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Filter bank of the wavelet."""
+        return self
+
+    @classmethod
+    def from_wavelet(cls, wavelet: Wavelet, dtype: torch.dtype) -> WaveletTensorTuple:
+        """Construct Wavelet named tuple from wavelet protocol member."""
+        return cls(
+            torch.tensor(wavelet.dec_lo, dtype=dtype),
+            torch.tensor(wavelet.dec_hi, dtype=dtype),
+            torch.tensor(wavelet.rec_lo, dtype=dtype),
+            torch.tensor(wavelet.rec_hi, dtype=dtype),
+        )
+
+
 def _as_wavelet(wavelet: Union[Wavelet, str]) -> Wavelet:
     """Ensure the input argument to be a pywt wavelet compatible object.
 

tests/test_jit.py

@@ -1,6 +1,6 @@
 """Ensure pytorch's torch.jit.trace feature works properly."""
 
-from typing import NamedTuple, Optional, Union
+from typing import Optional, Union
 
 import numpy as np
 import pytest
@@ -13,24 +13,6 @@
 from tests._mackey_glass import MackeyGenerator
 
 
-class WaveletTuple(NamedTuple):
-    """Replaces namedtuple("Wavelet", ("dec_lo", "dec_hi", "rec_lo", "rec_hi"))."""
-
-    dec_lo: torch.Tensor
-    dec_hi: torch.Tensor
-    rec_lo: torch.Tensor
-    rec_hi: torch.Tensor
-
-
-def _set_up_wavelet_tuple(wavelet: WaveletTuple, dtype: torch.dtype) -> WaveletTuple:
-    return WaveletTuple(
-        torch.tensor(wavelet.dec_lo).type(dtype),
-        torch.tensor(wavelet.dec_hi).type(dtype),
-        torch.tensor(wavelet.rec_lo).type(dtype),
-        torch.tensor(wavelet.rec_hi).type(dtype),
-    )
-
-
 def _to_jit_wavedec_fun(
     data: torch.Tensor, wavelet: Union[ptwt.Wavelet, str], level: Optional[int]
 ) -> list[torch.Tensor]:
@@ -53,7 +35,7 @@ def test_conv_fwt_jit(
 
     mackey_data_1 = torch.squeeze(generator(), -1).type(dtype)
     wavelet = pywt.Wavelet(wavelet_string)
-    wavelet = _set_up_wavelet_tuple(wavelet, dtype)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(wavelet, dtype)
 
     with pytest.warns(Warning):
         jit_wavedec = torch.jit.trace(  # type: ignore
@@ -105,7 +87,7 @@ def test_conv_fwt_jit_2d() -> None:
     rec = _to_jit_waverec_2(coeff, wavelet)
     assert np.allclose(rec.squeeze(1).numpy(), data.numpy())
 
-    wavelet = _set_up_wavelet_tuple(wavelet, dtype=torch.float64)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(wavelet, dtype=torch.float64)
     with pytest.warns(Warning):
         jit_wavedec2 = torch.jit.trace(  # type: ignore
             _to_jit_wavedec_2,
@@ -159,7 +141,7 @@ def test_conv_fwt_jit_3d() -> None:
     rec = _to_jit_waverec_3(coeff, wavelet)
     assert np.allclose(rec.squeeze(1).numpy(), data.numpy())
 
-    wavelet = _set_up_wavelet_tuple(wavelet, dtype=torch.float64)
+    wavelet = ptwt.WaveletTensorTuple.from_wavelet(wavelet, dtype=torch.float64)
     with pytest.warns(Warning):
         jit_wavedec3 = torch.jit.trace(  # type: ignore
             _to_jit_wavedec_3,