Merge pull request #2645 from ericpre/fix_rechunk_change_dtype

Use 'dask_auto' when rechunk=True in change_dtype.

CHANGES.rst

@@ -8,15 +8,16 @@ What's new
 RELEASE_next_patch (Unreleased)
 +++++++++++++++++++++++++++++++
 
-* Widgets plotting improvement and add `pick_tolerance` to plot preferences (`#2615 <https://github.com/hyperspy/hyperspy/pull/2615>`_)
+* Widgets plotting improvement and add ``pick_tolerance`` to plot preferences (`#2615 <https://github.com/hyperspy/hyperspy/pull/2615>`_)
 * Update external links in the loading data section of the user guide (`#2627 <https://github.com/hyperspy/hyperspy/pull/2627>`_)
 * Pass keyword argument to the image IO plugins (`#2627 <https://github.com/hyperspy/hyperspy/pull/2627>`_)
 * Drop support for numpy<1.16, in line with NEP 29 and fix protochip reader for numpy 1.20 (`#2616 <https://github.com/hyperspy/hyperspy/pull/2616>`_)
 * Run test suite against upstream dependencies (numpy, scipy, scikit-learn and scikit-image) (`#2616 <https://github.com/hyperspy/hyperspy/pull/2616>`_)
 * Improve error message when file not found (`#2597 <https://github.com/hyperspy/hyperspy/pull/2597>`_)
 * Add update instructions to user guide (`#2621 <https://github.com/hyperspy/hyperspy/pull/2621>`_)
 * Fix disconnect event when closing navigator only plot (fixes `#996 <https://github.com/hyperspy/hyperspy/issues/996>`_), (`#2631 <https://github.com/hyperspy/hyperspy/pull/2631>`_)
-* Improve plotting navigator of lazy signals, add `navigator` setter to lazy signals (`#2631 <https://github.com/hyperspy/hyperspy/pull/2631>`_)
+* Improve plotting navigator of lazy signals, add ``navigator`` setter to lazy signals (`#2631 <https://github.com/hyperspy/hyperspy/pull/2631>`_)
+* Use 'dask_auto' when rechunk=True in ``change_dtype`` for lazy signal (`#2645 <https://github.com/hyperspy/hyperspy/pull/2645>`_)
 * Use dask chunking when saving lazy signal instead of rechunking and leave the user to decide what is the suitable chunking (`#2629 <https://github.com/hyperspy/hyperspy/pull/2629>`_)
 * Fix incorrect chunksize when saving EMD NCEM file and specifying chunks (`#2629 <https://github.com/hyperspy/hyperspy/pull/2629>`_)
 

hyperspy/_signals/lazy.py

@@ -229,12 +229,16 @@ def _make_lazy(self, axis=None, rechunk=False, dtype=None):
         self.data = self._lazy_data(axis=axis, rechunk=rechunk, dtype=dtype)
 
     def change_dtype(self, dtype, rechunk=True):
+        # To be consistent with the rechunk argument of other method, we use
+        # 'dask_auto' in favour of a chunking which doesn't split signal space.
+        if rechunk:
+            rechunk = 'dask_auto'
         from hyperspy.misc import rgb_tools
         if not isinstance(dtype, np.dtype) and (dtype not in
                                                 rgb_tools.rgb_dtypes):
             dtype = np.dtype(dtype)
-            self._make_lazy(rechunk=rechunk, dtype=dtype)
         super().change_dtype(dtype)
+        self._make_lazy(rechunk=rechunk, dtype=dtype)
     change_dtype.__doc__ = BaseSignal.change_dtype.__doc__
 
     def _lazy_data(self, axis=None, rechunk=True, dtype=None):

hyperspy/tests/signal/test_lazy_tools.py

@@ -22,22 +22,26 @@
 from hyperspy.signals import Signal1D, Signal2D
 
 
-def test_lazy_changetype_rechunk():
-    ar = da.ones((50, 50, 256, 256), chunks=(5, 5, 256, 256), dtype="uint8")
+def test_lazy_changetype_rechunk_default():
+    ar = da.ones((50, 50, 512, 512), chunks=(5, 5, 128, 128), dtype="uint8")
     s = Signal2D(ar).as_lazy()
     s._make_lazy(rechunk=True)
     assert s.data.dtype is np.dtype("uint8")
     chunks_old = s.data.chunks
     s.change_dtype("float")
     assert s.data.dtype is np.dtype("float")
     chunks_new = s.data.chunks
+    # We expect more chunks
     assert len(chunks_old[0]) * len(chunks_old[1]) < len(chunks_new[0]) * len(
         chunks_new[1]
     )
     s.change_dtype("uint8")
     assert s.data.dtype is np.dtype("uint8")
     chunks_newest = s.data.chunks
-    assert chunks_newest == chunks_new
+    # We expect less chunks
+    assert len(chunks_newest[0]) * len(chunks_newest[1]) < len(chunks_new[0]) * len(
+        chunks_new[1]
+    )
 
 
 def test_lazy_changetype_rechunk_False():