fix types.

Author: v0lta

src/ptwt/continuous_transform.py

@@ -196,7 +196,7 @@ def _integrate(
         if type(arr) is np.ndarray:
             integral = np.cumsum(arr)
         elif type(arr) is torch.Tensor:
-            integral = torch.cumsum(arr, -1)
+            integral = torch.cumsum(arr, -1) # type: ignore
         else:
             raise TypeError("Only ndarrays or tensors are integratable.")
         integral *= step
@@ -272,8 +272,8 @@ def wavefun(
         """Define a grid and evaluate the wavelet on it."""
         length = 2**precision
         # load the bounds from untyped pywt code.
-        lower_bound: float = float(self.lower_bound)
-        upper_bound: float = float(self.upper_bound)
+        lower_bound: float = float(self.lower_bound)  # type: ignore
+        upper_bound: float = float(self.upper_bound)  # type: ignore
         grid = torch.linspace(
             lower_bound,
             upper_bound,
@@ -292,10 +292,10 @@ def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
         shannon = (
             torch.sqrt(self.bandwidth)
             * (
-                torch.sin(torch.pi * self.bandwidth * grid_values)  # type: ignore
+                torch.sin(torch.pi * self.bandwidth * grid_values)
                 / (torch.pi * self.bandwidth * grid_values)
             )
-            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)  # type: ignore
+            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
         )
         return shannon
 
@@ -307,8 +307,8 @@ def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
         """Return numerical values for the wavelet on a grid."""
         morlet = (
             1.0
-            / torch.sqrt(torch.pi * self.bandwidth)  # type: ignore
+            / torch.sqrt(torch.pi * self.bandwidth)
             * torch.exp(-(grid_values**2) / self.bandwidth)
-            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)  # type: ignore
+            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
         )
         return morlet

src/ptwt/conv_transform_2.py

@@ -216,7 +216,7 @@ def wavedec2(
         result_lst = _map_result(result_lst, _unfold_axes2)
 
     if axes != (-2, -1):
-        undo_swap_fn = partial(_undo_swap_axes, axes=axes)
+        undo_swap_fn = partial(_undo_swap_axes, axes=list(axes))
         result_lst = _map_result(result_lst, undo_swap_fn)
 
     return result_lst

src/ptwt/conv_transform_3.py

@@ -205,7 +205,7 @@ def wavedec3(
         result_lst = _map_result(result_lst, _unfold_axes_fn)
 
     if tuple(axes) != (-3, -2, -1):
-        undo_swap_fn = partial(_undo_swap_axes, axes=axes)
+        undo_swap_fn = partial(_undo_swap_axes, axes=list(axes))
         result_lst = _map_result(result_lst, undo_swap_fn)
 
     return result_lst

src/ptwt/matmul_transform_2.py

@@ -546,7 +546,7 @@ def __call__(
             split_list = _map_result(split_list, _unfold_axes2)
 
         if self.axes != (-2, -1):
-            undo_swap_fn = partial(_undo_swap_axes, axes=self.axes)
+            undo_swap_fn = partial(_undo_swap_axes, axes=list(self.axes))
             split_list = _map_result(split_list, undo_swap_fn)
 
         return split_list[::-1]

src/ptwt/matmul_transform_3.py

@@ -139,7 +139,7 @@ def _construct_analysis_matrices(
             matrix_construction_fun = partial(
                 construct_boundary_a,
                 wavelet=self.wavelet,
-                boundary=self.boundary,
+                boundary=self.boundary,  # type: ignore
                 device=device,
                 dtype=dtype,
             )
@@ -267,7 +267,7 @@ def _split_rec(
             split_list = _map_result(split_list, _unfold_axes_fn)
 
         if self.axes != (-3, -2, -1):
-            undo_swap_fn = partial(_undo_swap_axes, axes=self.axes)
+            undo_swap_fn = partial(_undo_swap_axes, axes=list(self.axes))
             split_list = _map_result(split_list, undo_swap_fn)
 
         return split_list[::-1]

src/ptwt/packets.py

@@ -105,7 +105,7 @@ def __init__(
             if len(data.shape) == 1:
                 # add a batch dimension.
                 data = data.unsqueeze(0)
-            self.transform(data, maxlevel)  # type: ignore
+            self.transform(data, maxlevel)
         else:
             self.data = {}
 
@@ -177,7 +177,7 @@ def _get_wavedec(
             return self._matrix_wavedec_dict[length]
         else:
             return partial(
-                wavedec, wavelet=self.wavelet, level=1, mode=self.mode, axis=self.axis
+                wavedec, wavelet=self.wavelet, level=1, mode=self.mode, axis=self.axis # type: ignore
             )
 
     def _get_waverec(

src/ptwt/separable_conv_transform.py

@@ -165,9 +165,9 @@ def _separable_conv_waverecn(
 
     approx: torch.Tensor = coeffs[0]
     for level_dict in coeffs[1:]:
-        keys = list(level_dict.keys())
-        level_dict["a" * max(map(len, keys))] = approx
-        approx = _separable_conv_idwtn(level_dict, wavelet)
+        keys = list(level_dict.keys()) # type: ignore
+        level_dict["a" * max(map(len, keys))] = approx  # type: ignore
+        approx = _separable_conv_idwtn(level_dict, wavelet)  # type: ignore
     return approx
 
 
@@ -236,7 +236,7 @@ def fswavedec2(
         res = _map_result(res, _unfold_axes2)
 
     if axes != (-2, -1):
-        undo_swap_fn = partial(_undo_swap_axes, axes=axes)
+        undo_swap_fn = partial(_undo_swap_axes, axes=list(axes))
         res = _map_result(res, undo_swap_fn)
 
     return res
@@ -308,7 +308,7 @@ def fswavedec3(
         res = _map_result(res, _unfold_axes3)
 
     if axes != (-3, -2, -1):
-        undo_swap_fn = partial(_undo_swap_axes, axes=axes)
+        undo_swap_fn = partial(_undo_swap_axes, axes=list(axes))
         res = _map_result(res, undo_swap_fn)
 
     return res