Annotation as NewType dtcwt.py

Author: cako

pylops/signalprocessing/dtcwt.py

@@ -1,6 +1,6 @@
 __all__ = ["DTCWT"]
 
-from typing import Any, Union
+from typing import Any, NewType, Union
 
 import numpy as np
 
@@ -19,6 +19,8 @@
 else:
     pyramid_type = Any
 
+PyramidType = NewType("PyramidType", pyramid_type)
+
 
 class DTCWT(LinearOperator):
     r"""Dual-Tree Complex Wavelet Transform
@@ -146,16 +148,12 @@ def _nd_to_2d(self, arr_nd: NDArray) -> NDArray:
         arr_2d = arr_nd.reshape(self.dims[self.axis], -1).squeeze()
         return arr_2d
 
-    def _coeff_to_array(
-        self, pyr: pyramid_type
-    ) -> NDArray:  # cannot use dtcwt types as it may not be installed
+    def _coeff_to_array(self, pyr: PyramidType) -> NDArray:
         highpass_coeffs = np.vstack([h for h in pyr.highpasses])
         coeffs = np.concatenate((highpass_coeffs, pyr.lowpass), axis=0)
         return coeffs
 
-    def _array_to_coeff(
-        self, X: NDArray
-    ) -> pyramid_type:  # cannot use dtcwt types as it may not be installed
+    def _array_to_coeff(self, X: NDArray) -> PyramidType:
         lowpass = (X[-self.lowpass_size :].real).reshape((-1, self.otherdims))
         _ptr = 0
         highpasses = ()
@@ -166,9 +164,7 @@ def _array_to_coeff(
             highpasses += (_h,)
         return dtcwt.Pyramid(lowpass, highpasses)
 
-    def get_pyramid(
-        self, x: NDArray
-    ) -> pyramid_type:  # cannot use dtcwt types as it may not be installed
+    def get_pyramid(self, x: NDArray) -> PyramidType:
         """Return Pyramid object from flat real-valued array"""
         return self._array_to_coeff(x[0] + 1j * x[1])