Merge pull request #2754 from magnunor/fix_lazy_map_nav_chunksize_1

Add fix for lazy s.map and nav chunks with size of 1

hyperspy/misc/utils.py

@@ -1229,10 +1229,11 @@ def process_function_blockwise(data,
             output_array[islice] = function(data[islice],
                                             **iter_dict,
                                             **kwargs)
-    try:
-        output_array = output_array.squeeze(-1)
-    except ValueError:
-        pass
+    if not (chunk_nav_shape == output_array.shape):
+        try:
+            output_array = output_array.squeeze(-1)
+        except ValueError:
+            pass
     return output_array
 
 

hyperspy/tests/signals/test_map_method.py

@@ -420,3 +420,21 @@ def test_map_ufunc(caplog):
     assert np.log(s) == s.map(np.log)
     np.testing.assert_allclose(s.data, np.log(data))
     assert "can direcly operate on hyperspy signals" in caplog.records[0].message
+
+
+class TestLazyNavChunkSize1:
+    @staticmethod
+    def afunction(input_data):
+        return np.array([1, 2, 3])
+
+    def test_signal2d(self):
+        dask_array = da.zeros((10, 15, 32, 32), chunks=(1, 1, 32, 32))
+        s = hs.signals.Signal2D(dask_array).as_lazy()
+        s_out = s.map(self.afunction, inplace=False, parallel=False, ragged=True)
+        s_out.compute()
+
+    def test_signal1d(self):
+        dask_array = da.zeros((10, 15, 32), chunks=(1, 1, 32))
+        s = hs.signals.Signal1D(dask_array).as_lazy()
+        s_out = s.map(self.afunction, inplace=False, parallel=False, ragged=True)
+        s_out.compute()