diff --git a/pylops/signalprocessing/dtcwt.py b/pylops/signalprocessing/dtcwt.py index da149585..75b19f8d 100644 --- a/pylops/signalprocessing/dtcwt.py +++ b/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])