diff --git a/src/ptwt/_util.py b/src/ptwt/_util.py index 52e53f68..ae589761 100644 --- a/src/ptwt/_util.py +++ b/src/ptwt/_util.py @@ -850,3 +850,87 @@ def _get_default_axes(n: int) -> tuple[int, ...]: if n < 1: raise ValueError(f"only natural number dimensions are allowed. given: {n}") return tuple(range(-n, 0)) + + +def _preprocess_deconstruction( + data: torch.Tensor, + wavelet: Union[Wavelet, str], + *, + ndim: int, + axes: AxisHint = None, +) -> tuple[torch.Tensor, list[int], torch.Tensor, torch.Tensor, torch.Tensor]: + data, ds = _preprocess_tensor(data, ndim=ndim, axes=axes) + dec_lo, dec_hi, _, _ = _get_filter_tensors( + wavelet, flip=True, device=data.device, dtype=data.dtype + ) + dec_filt = _construct_nd_filt(dec_lo, dec_hi, n=ndim) + return data, ds, dec_lo, dec_hi, dec_filt + + +def _construct_nd_filt(lo: torch.Tensor, hi: torch.Tensor, n: int) -> torch.Tensor: + if n == 1: + return _construct_1d_filt(lo, hi) + elif n == 2: + return _construct_2d_filt(lo, hi) + elif n == 3: + return _construct_3d_filt(lo, hi) + else: + raise NotImplementedError() + + +def _construct_1d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: + """Construct one-dimensional filters.""" + return torch.stack([lo, hi], 0) + + +def _construct_2d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: + """Construct two-dimensional filters using outer products. + + Args: + lo (torch.Tensor): Low-pass input filter. + hi (torch.Tensor): High-pass input filter + + Returns: + Stacked 2d-filters of dimension + + [2^2, 1, height, width]. + + The four filters are ordered ll, lh, hl, hh. + + """ + ll = _outer(lo, lo) + lh = _outer(hi, lo) + hl = _outer(lo, hi) + hh = _outer(hi, hi) + filt = torch.stack([ll, lh, hl, hh], 0) + filt = filt.unsqueeze(1) + return filt + + +def _construct_3d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: + """Construct three-dimensional filters using outer products. + + Args: + lo (torch.Tensor): Low-pass input filter. + hi (torch.Tensor): High-pass input filter + + Returns: + Stacked 3d filters of dimension:: + + [2^3, 1, length, height, width]. + + The four filters are ordered ll, lh, hl, hh. + """ + dim_size = lo.shape[-1] + size = [dim_size] * 3 + lll = _outer(lo, _outer(lo, lo)).reshape(size) + llh = _outer(lo, _outer(lo, hi)).reshape(size) + lhl = _outer(lo, _outer(hi, lo)).reshape(size) + lhh = _outer(lo, _outer(hi, hi)).reshape(size) + hll = _outer(hi, _outer(lo, lo)).reshape(size) + hlh = _outer(hi, _outer(lo, hi)).reshape(size) + hhl = _outer(hi, _outer(hi, lo)).reshape(size) + hhh = _outer(hi, _outer(hi, hi)).reshape(size) + filt = torch.stack([lll, llh, lhl, lhh, hll, hlh, hhl, hhh], 0) + filt = filt.unsqueeze(1) + return filt diff --git a/src/ptwt/conv_transform.py b/src/ptwt/conv_transform.py index b5c39f10..22c954b5 100644 --- a/src/ptwt/conv_transform.py +++ b/src/ptwt/conv_transform.py @@ -22,7 +22,7 @@ _postprocess_coeffs, _postprocess_tensor, _preprocess_coeffs, - _preprocess_tensor, + _preprocess_deconstruction, _translate_boundary_strings, ) from .constants import BoundaryMode, Wavelet, WaveletCoeff1d @@ -91,7 +91,7 @@ def wavedec( Args: data (torch.Tensor): The input time series to transform. - By default the last axis is transformed. + By default, the last axis is transformed. wavelet (Wavelet or str): A pywt wavelet compatible object or the name of a pywt wavelet. Please consider the output from ``pywt.wavelist(kind='discrete')`` @@ -122,22 +122,19 @@ def wavedec( >>> # compute the forward fwt coefficients >>> ptwt.wavedec(data, 'haar', mode='zero', level=2) """ - data, ds = _preprocess_tensor(data, ndim=1, axes=axis) - - dec_lo, dec_hi, _, _ = _get_filter_tensors( - wavelet, flip=True, device=data.device, dtype=data.dtype + data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction( + data, wavelet, axes=axis, ndim=1 ) - filt_len = dec_lo.shape[-1] - filt = torch.stack([dec_lo, dec_hi], 0) if level is None: + filt_len = dec_lo.shape[-1] level = pywt.dwt_max_level(data.shape[-1], filt_len) result_list = [] res_lo = data for _ in range(level): res_lo = _fwt_pad(res_lo, wavelet, mode=mode) - res = torch.nn.functional.conv1d(res_lo, filt, stride=2) + res = torch.nn.functional.conv1d(res_lo, dec_filt, stride=2) res_lo, res_hi = torch.split(res, 1, 1) result_list.append(res_hi.squeeze(1)) result_list.append(res_lo.squeeze(1)) diff --git a/src/ptwt/conv_transform_2.py b/src/ptwt/conv_transform_2.py index 8d14d40e..a0ff6804 100644 --- a/src/ptwt/conv_transform_2.py +++ b/src/ptwt/conv_transform_2.py @@ -15,15 +15,15 @@ AxisHint, _adjust_padding_at_reconstruction, _check_same_device_dtype, + _construct_2d_filt, _get_filter_tensors, _get_padding_n, _group_for_symmetric, - _outer, _pad_symmetric, _postprocess_coeffs, _postprocess_tensor, _preprocess_coeffs, - _preprocess_tensor, + _preprocess_deconstruction, _translate_boundary_strings, ) from .constants import BoundaryMode, Wavelet, WaveletCoeff2d, WaveletDetailTuple2d @@ -31,30 +31,6 @@ __all__ = ["wavedec2", "waverec2"] -def _construct_2d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: - """Construct two-dimensional filters using outer products. - - Args: - lo (torch.Tensor): Low-pass input filter. - hi (torch.Tensor): High-pass input filter - - Returns: - Stacked 2d-filters of dimension - - [2^2, 1, height, width]. - - The four filters are ordered ll, lh, hl, hh. - - """ - ll = _outer(lo, lo) - lh = _outer(hi, lo) - hl = _outer(lo, hi) - hh = _outer(hi, hi) - filt = torch.stack([ll, lh, hl, hh], 0) - filt = filt.unsqueeze(1) - return filt - - def _fwt_pad2( data: torch.Tensor, wavelet: Union[Wavelet, str], @@ -154,11 +130,9 @@ def wavedec2( >>> coefficients = ptwt.wavedec2(data, "haar", level=2, mode="zero") """ - data, ds = _preprocess_tensor(data, ndim=2, axes=axes) - dec_lo, dec_hi, _, _ = _get_filter_tensors( - wavelet, flip=True, device=data.device, dtype=data.dtype + data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction( + data, wavelet, axes=axes, ndim=2 ) - dec_filt = _construct_2d_filt(lo=dec_lo, hi=dec_hi) if level is None: level = pywt.dwtn_max_level([data.shape[-1], data.shape[-2]], wavelet) diff --git a/src/ptwt/conv_transform_3.py b/src/ptwt/conv_transform_3.py index fac62c9e..f6249a86 100644 --- a/src/ptwt/conv_transform_3.py +++ b/src/ptwt/conv_transform_3.py @@ -15,15 +15,15 @@ _adjust_padding_at_reconstruction, _as_wavelet, _check_same_device_dtype, + _construct_3d_filt, _get_filter_tensors, _get_padding_n, _group_for_symmetric, - _outer, _pad_symmetric, _postprocess_coeffs, _postprocess_tensor, _preprocess_coeffs, - _preprocess_tensor, + _preprocess_deconstruction, _translate_boundary_strings, ) from .constants import BoundaryMode, Wavelet, WaveletCoeffNd, WaveletDetailDict @@ -31,35 +31,6 @@ __all__ = ["wavedec3", "waverec3"] -def _construct_3d_filt(lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: - """Construct three-dimensional filters using outer products. - - Args: - lo (torch.Tensor): Low-pass input filter. - hi (torch.Tensor): High-pass input filter - - Returns: - Stacked 3d filters of dimension:: - - [2^3, 1, length, height, width]. - - The four filters are ordered ll, lh, hl, hh. - """ - dim_size = lo.shape[-1] - size = [dim_size] * 3 - lll = _outer(lo, _outer(lo, lo)).reshape(size) - llh = _outer(lo, _outer(lo, hi)).reshape(size) - lhl = _outer(lo, _outer(hi, lo)).reshape(size) - lhh = _outer(lo, _outer(hi, hi)).reshape(size) - hll = _outer(hi, _outer(lo, lo)).reshape(size) - hlh = _outer(hi, _outer(lo, hi)).reshape(size) - hhl = _outer(hi, _outer(hi, lo)).reshape(size) - hhh = _outer(hi, _outer(hi, hi)).reshape(size) - filt = torch.stack([lll, llh, lhl, lhh, hll, hlh, hhl, hhh], 0) - filt = filt.unsqueeze(1) - return filt - - def _fwt_pad3( data: torch.Tensor, wavelet: Union[Wavelet, str], @@ -137,13 +108,10 @@ def wavedec3( >>> data = torch.randn(5, 16, 16, 16) >>> transformed = ptwt.wavedec3(data, "haar", level=2, mode="reflect") """ - data, ds = _preprocess_tensor(data, ndim=3, axes=axes) - wavelet = _as_wavelet(wavelet) - dec_lo, dec_hi, _, _ = _get_filter_tensors( - wavelet, flip=True, device=data.device, dtype=data.dtype + data, ds, dec_lo, dec_hi, dec_filt = _preprocess_deconstruction( + data, wavelet, axes=axes, ndim=3 ) - dec_filt = _construct_3d_filt(lo=dec_lo, hi=dec_hi) if level is None: level = pywt.dwtn_max_level( diff --git a/src/ptwt/matmul_transform_2.py b/src/ptwt/matmul_transform_2.py index d8d4c1ec..02ab458b 100644 --- a/src/ptwt/matmul_transform_2.py +++ b/src/ptwt/matmul_transform_2.py @@ -15,6 +15,7 @@ AxisHint, _as_wavelet, _check_same_device_dtype, + _construct_2d_filt, _deprecated_alias, _ensure_axes, _get_filter_tensors, @@ -32,7 +33,7 @@ WaveletCoeff2d, WaveletDetailTuple2d, ) -from .conv_transform_2 import _construct_2d_filt, _fwt_pad2 +from .conv_transform_2 import _fwt_pad2 from .matmul_transform import ( BaseMatrixWaveDec, construct_boundary_a,