BUG: infer multiresolution levels from transformed axes

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)

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)