Merge pull request #860 from rgommers/waveletpacketnd-followup

BUG: fix indexing in `reconstruct` method of 1-D/2-D wavelet packet

Author: rgommers

pywt/_wavelet_packets.py

@@ -675,6 +675,9 @@ def _reconstruct(self, update):
             )
         else:
             rec = idwtn(coeffs, self.wavelet, self.mode, axes=self.axes)
+            if self._data_shape is not None and (
+                    rec.shape != self._data_shape):
+                rec = rec[tuple([slice(sz) for sz in self._data_shape])]
             if update:
                 self.data = rec
             return rec
@@ -741,7 +744,7 @@ def reconstruct(self, update=True):
         if self.has_any_subnode:
             data = super().reconstruct(update)
             if self.data_size is not None and (data.shape != self.data_size):
-                data = data[[slice(sz) for sz in self.data_size]]
+                data = data[tuple(slice(sz) for sz in self.data_size)]
             if update:
                 self.data = data
             return data
@@ -867,7 +870,7 @@ def reconstruct(self, update=True):
         if self.has_any_subnode:
             data = super().reconstruct(update)
             if self.data_size is not None and (data.shape != self.data_size):
-                data = data[[slice(sz) for sz in self.data_size]]
+                data = data[tuple(slice(sz) for sz in self.data_size)]
             if update:
                 self.data = data
             return data

pywt/tests/test_wpnd.py

@@ -168,3 +168,56 @@ def test_wavelet_packet_axes():
     # must have non-duplicate axes
     assert_raises(ValueError, pywt.WaveletPacketND, data=x, wavelet='db1',
                   axes=(0, 0))
+
+
+def test_wavelet_packet_odd_shape():
+    # the reconstruction has to be trimmed back to the original shape when a
+    # transformed axis has an odd length
+    x = np.arange(2 * 2 * 3, dtype=np.float64).reshape(2, 2, 3)
+    wp = pywt.WaveletPacketND(data=x, wavelet='haar', mode='symmetric',
+                              axes=(0, 1, 2))
+    wp.get_level(1)
+    r = wp.reconstruct(update=False)
+    assert_equal(r.shape, x.shape)
+    assert_allclose(r, x, atol=1e-12, rtol=1e-12)
+
+    rstate = np.random.RandomState(0)
+    y = rstate.standard_normal((9, 11))
+    for level in [1, 2, 3]:
+        wp = pywt.WaveletPacketND(data=y, wavelet='haar', mode='symmetric')
+        wp.get_level(level)
+        r = wp.reconstruct(update=False)
+        assert_equal(r.shape, y.shape)
+        assert_allclose(r, y, rtol=1e-12)
+
+    # only some of the axes transformed
+    z = rstate.standard_normal((3, 5, 7))
+    wp = pywt.WaveletPacketND(data=z, wavelet='haar', mode='symmetric',
+                              axes=(1, 2))
+    wp.get_level(1)
+    assert_allclose(wp.reconstruct(update=False), z, rtol=1e-12)
+
+    # reconstructing again after an update gives the same result
+    wp = pywt.WaveletPacketND(data=y, wavelet='haar', mode='symmetric')
+    wp.get_level(2)
+    assert_allclose(wp.reconstruct(update=True), y, rtol=1e-12)
+    assert_allclose(wp.reconstruct(update=True), y, rtol=1e-12)
+
+
+def test_wavelet_packet_odd_shape_subnode():
+    # a subnode reconstructs to the shape of its own coefficients, so an
+    # update does not grow the data stored in the node
+    rstate = np.random.RandomState(0)
+    y = rstate.standard_normal((9, 11))
+    wp = pywt.WaveletPacketND(data=y, wavelet='haar', mode='symmetric')
+    wp.get_level(2)
+
+    shape = wp['aa'].data.shape
+    assert_equal(wp['aa'].reconstruct(update=False).shape, shape)
+    wp['aa'].reconstruct(update=True)
+    assert_equal(wp['aa'].data.shape, shape)
+
+    # WaveletPacket2D already behaves this way
+    wp2 = pywt.WaveletPacket2D(data=y, wavelet='haar', mode='symmetric')
+    wp2.get_level(2)
+    assert_equal(wp2['a'].reconstruct(update=False).shape, shape)