Improve plotting of lazy signals, by keeping current chunk

hyperspy/_signals/lazy.py

@@ -162,8 +162,10 @@ def close_file(self):
                 _logger.exception("Failed to close lazy Signal file")
 
     def _clear_temp_dask_data(self, obj=None):
-        del self._temp_dask_chunk
-        del self._temp_dask_chunk_slice
+        if hasattr(self, '_temp_dask_chunk'):
+            del self._temp_dask_chunk
+        if hasattr(self, '_temp_dask_chunk_slice'):
+            del self._temp_dask_chunk_slice
 
     def _get_dask_chunks(self, axis=None, dtype=None):
         """Returns dask chunks.

hyperspy/tests/signal/test_lazy.py

@@ -303,11 +303,11 @@ def test_changed_data(self):
         s = _lazy_signals.LazySignal2D(da.zeros((6, 6, 8, 8), chunks=(2, 2, 4, 4)))
         position = s.axes_manager._getitem_tuple
         s._get_temporary_dask_chunk(position)
-        assert hasattr(s._temp_data_chunk)
-        assert hasattr(s._temp_data_chunk_slice)
+        assert hasattr(s, "_temp_dask_chunk")
+        assert hasattr(s, "_temp_dask_chunk_slice")
         s.events.data_changed.trigger(None)
-        assert not hasattr(s._temp_data_chunk)
-        assert not hasattr(s._temp_data_chunk_slice)
+        assert not hasattr(s, "_temp_dask_chunk")
+        assert not hasattr(s, "_temp_dask_chunk_slice")
 
     def test_map(self):
         s = _lazy_signals.LazySignal2D(da.zeros((6, 6, 8, 8), chunks=(2, 2, 4, 4)))
@@ -317,6 +317,14 @@ def test_map(self):
         s.__call__()
         assert len(s._temp_dask_chunk.shape) == 3
 
+    def test_clear_temp_dask_data(self):
+        s = _lazy_signals.LazySignal2D(da.zeros((6, 6, 8, 8), chunks=(2, 2, 4, 4)))
+        s.__call__()
+        s._clear_temp_dask_data(self)
+        assert not hasattr(s, "_temp_dask_data")
+        assert not hasattr(s, "_temp_dask_data_slice")
+        s._clear_temp_dask_data(self)
+
 
 class TestLazyPlot:
     def test_correct_value(self):