diff --git a/pywt/_mra.py b/pywt/_mra.py index a41431e3..5552bfa1 100644 --- a/pywt/_mra.py +++ b/pywt/_mra.py @@ -161,11 +161,11 @@ def mra2(data, wavelet, level=None, axes=(-2, -1), transform='swt2', Wavelet to use. This can also be a tuple containing a wavelet to apply along each axis in `axes`. level : int, optional - Decomposition level (must be >= 0). If level is None (default) then it - will be calculated using the `dwt_max_level` function. + Decomposition level (must be >= 0). If level is None (default), it is + calculated for the transformed axes using `swt_max_level` for SWT or + `dwt_max_level` for DWT. axes : 2-tuple of ints, optional Axes over which to compute the DWT. Repeated elements are not allowed. - Currently only available when ``transform='dwt2'``. transform : {'dwt2', 'swt2'} Whether to use the DWT or SWT for the transforms. mode : str or 2-tuple of str, optional @@ -207,7 +207,8 @@ def mra2(data, wavelet, level=None, axes=(-2, -1), transform='swt2', raise ValueError( "transform swt only supports mode='periodization'") if level is None: - level = min(swt_max_level(s) for s in data.shape) + _, axes_shapes, _ = _prep_axes_wavedecn(data.shape, axes) + level = min(swt_max_level(s) for s in axes_shapes) kwargs = {"wavelet": wavelet, "axes": axes, "norm": True} forward = partial(swt2, level=level, trim_approx=True, **kwargs) inverse = partial(iswt2, **kwargs) @@ -301,8 +302,9 @@ def mran(data, wavelet, level=None, axes=None, transform='swtn', Wavelet to use. This can also be a tuple containing a wavelet to apply along each axis in `axes`. level : int, optional - Decomposition level (must be >= 0). If level is None (default) then it - will be calculated using the `dwt_max_level` function. + Decomposition level (must be >= 0). If level is None (default), it is + calculated for the transformed axes using `swt_max_level` for SWT or + `dwt_max_level` for DWT. axes : tuple of ints, optional Axes over which to compute the DWT. Repeated elements are not allowed. transform : {'dwtn', 'swtn'} @@ -350,7 +352,7 @@ def mran(data, wavelet, level=None, axes=None, transform='swtn', raise ValueError( "transform swt only supports mode='periodization'") if level is None: - level = min(swt_max_level(s) for s in data.shape) + level = min(swt_max_level(s) for s in axes_shapes) kwargs = {"wavelet": wavelets, "axes": axes, "norm": True} forward = partial(swtn, level=level, trim_approx=True, **kwargs) inverse = partial(iswtn, **kwargs) diff --git a/pywt/tests/test_mra_default_level.py b/pywt/tests/test_mra_default_level.py new file mode 100644 index 00000000..7c779fc5 --- /dev/null +++ b/pywt/tests/test_mra_default_level.py @@ -0,0 +1,29 @@ +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import pywt + + +@pytest.mark.parametrize('nd', [False, True]) +@pytest.mark.parametrize('batch_size', [1, 2, 3]) +@pytest.mark.parametrize('batch_axis', [0, 1, 2]) +def test_default_swt_mra_level_uses_only_transformed_axes(nd, batch_size, batch_axis): + image = np.arange(128, dtype=float).reshape(16, 8) + data = np.stack([image + i for i in range(batch_size)], axis=batch_axis) + axes = tuple(i for i in range(3) if i != batch_axis) + # Negative axes should select the same spatial dimensions. + axes = tuple(i - data.ndim for i in axes) + forward = pywt.mran if nd else pywt.mra2 + actual = forward(data, 'db1', axes=axes) + expected = forward(image, 'db1') + assert len(actual) == len(expected) == 4 + for i in range(batch_size): + single = forward(image + i, 'db1') + assert_allclose(np.take(actual[0], i, axis=batch_axis), single[0]) + for got_details, expected_details in zip(actual[1:], single[1:]): + if nd: + got_details = got_details.values() + expected_details = expected_details.values() + for got, expected_array in zip(got_details, expected_details): + assert_allclose(np.take(got, i, axis=batch_axis), expected_array)