Skip to content

Commit

Permalink
Add fix for lazy s.map and nav chunks with size of 1
Browse files Browse the repository at this point in the history
  • Loading branch information
magnunor committed May 25, 2021
1 parent 4d128c1 commit ea0f24c
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 4 deletions.
9 changes: 5 additions & 4 deletions hyperspy/misc/utils.py
Expand Up @@ -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


Expand Down
18 changes: 18 additions & 0 deletions hyperspy/tests/signals/test_map_method.py
Expand Up @@ -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()

0 comments on commit ea0f24c

Please sign in to comment.