From 02c80742e2a60d4d2f11756f2fff88fc95b14fcb Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 6 Jun 2024 10:48:13 +0200 Subject: [PATCH 1/5] Specify meta in da.map_blocks --- lib/iris/_lazy_data.py | 17 ++++++++++++----- lib/iris/experimental/ugrid/utils.py | 5 ++++- .../regrid/test_RectilinearRegridder.py | 3 +++ lib/iris/tests/unit/analysis/test_PERCENTILE.py | 10 ++++++---- .../unit/lazy_data/test_map_complete_blocks.py | 6 ++++++ 5 files changed, 31 insertions(+), 10 deletions(-) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index b1f8e9aa85..1e094d8ebe 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -535,9 +535,12 @@ def lazy_elementwise(lazy_array, elementwise_op): # This makes good practical sense for unit conversions, as a Unit.convert # call may cast to float, or not, depending on unit equality : Thus, it's # much safer to get udunits to decide that for us. - dtype = elementwise_op(np.zeros(1, lazy_array.dtype)).dtype + meta = da.utils.meta_from_array(lazy_array) + new_meta = elementwise_op(meta) - return da.map_blocks(elementwise_op, lazy_array, dtype=dtype) + return da.map_blocks( + elementwise_op, lazy_array, dtype=new_meta.dtype, meta=new_meta + ) def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): @@ -595,8 +598,12 @@ def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): for dim, size in zip(dims, out_sizes): out_chunks[dim] = size - result = data.map_blocks( - func, *args, chunks=out_chunks, dtype=src.dtype, **kwargs - ) + # Assume operation does not change dtype and meta if not specified. + if "meta" not in kwargs: + kwargs["meta"] = da.utils.meta_from_array(data) + if "dtype" not in kwargs: + kwargs["dtype"] = kwargs["meta"].dtype + + result = data.map_blocks(func, *args, chunks=out_chunks, **kwargs) return result diff --git a/lib/iris/experimental/ugrid/utils.py b/lib/iris/experimental/ugrid/utils.py index dcf5462ad5..225e28e94e 100644 --- a/lib/iris/experimental/ugrid/utils.py +++ b/lib/iris/experimental/ugrid/utils.py @@ -274,6 +274,9 @@ def fill_region(target, regiondata, regioninds): # Notes on resultant calculation properties: # 1. map_blocks is chunk-mapped, so it is parallelisable and space-saving # 2. However, fetching less than a whole chunk is not efficient + meta = np.ma.array( + np.empty((0,) * result_array.ndim, dtype=result_array.dtype), mask=True + ) for cube in submesh_cubes: # Lazy data array from the region cube sub_data = cube.lazy_data() @@ -297,7 +300,7 @@ def fill_region(target, regiondata, regioninds): sub_data, indarr, dtype=result_array.dtype, - meta=np.ndarray, + meta=meta, ) # Construct the result cube diff --git a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py index 284d52d3f9..41cb107f37 100644 --- a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py +++ b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py @@ -481,6 +481,9 @@ def test_lazy_regrid(self): self.assertTrue(result.has_lazy_data()) expected = self.regridder(self.cube) self.assertTrue(result == expected) + self.assertTrue( + isinstance(da.utils.meta_from_array(result.core_data()), np.ma.MaskedArray) + ) class Test___call____invalid_types(tests.IrisTest): diff --git a/lib/iris/tests/unit/analysis/test_PERCENTILE.py b/lib/iris/tests/unit/analysis/test_PERCENTILE.py index 0d759c621f..72218af830 100644 --- a/lib/iris/tests/unit/analysis/test_PERCENTILE.py +++ b/lib/iris/tests/unit/analysis/test_PERCENTILE.py @@ -155,10 +155,10 @@ def test_default_kwargs_passed(self, mocked_mquantiles): if self.lazy: data = as_lazy_data(data) - self.agg_method(data, axis=axis, percent=percent) + result = self.agg_method(data, axis=axis, percent=percent) # Trigger calculation for lazy case. - as_concrete_data(data) + as_concrete_data(result) for key in ["alphap", "betap"]: self.assertEqual(mocked_mquantiles.call_args.kwargs[key], 1) @@ -170,10 +170,12 @@ def test_chosen_kwargs_passed(self, mocked_mquantiles): if self.lazy: data = as_lazy_data(data) - self.agg_method(data, axis=axis, percent=percent, alphap=0.6, betap=0.5) + result = self.agg_method( + data, axis=axis, percent=percent, alphap=0.6, betap=0.5 + ) # Trigger calculation for lazy case. - as_concrete_data(data) + as_concrete_data(result) for key, val in zip(["alphap", "betap"], [0.6, 0.5]): self.assertEqual(mocked_mquantiles.call_args.kwargs[key], val) diff --git a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py index 7403a5611e..c0f3a49f27 100644 --- a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py +++ b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py @@ -61,6 +61,12 @@ def test_dask_array_input(self): self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), self.func_result) + def test_dask_masked_array_input(self): + array = da.ma.masked_array(np.arange(2), mask=np.arange(2)) + result = map_complete_blocks(array, self.func, dims=tuple(), out_sizes=tuple()) + self.assertTrue(is_lazy_data(result)) + self.assertArrayEqual(result.compute(), np.ma.masked_array([1, 2], mask=[0, 1])) + def test_rechunk(self): lazy_array = da.asarray(self.array, chunks=((1, 1), (2, 2))) cube, _ = create_mock_cube(lazy_array) From 916408d4f9e26dcd15ed6b70e7c23fdbd11deb96 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 3 Jul 2024 09:29:58 +0200 Subject: [PATCH 2/5] Improve tests --- .../analysis/regrid/test_RectilinearRegridder.py | 10 ++++++++++ .../unit/lazy_data/test_map_complete_blocks.py | 15 ++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py index 41cb107f37..474ffb02cc 100644 --- a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py +++ b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py @@ -474,16 +474,26 @@ def setUp(self): self.args = ("linear", "mask") self.regridder = Regridder(self.cube, self.cube, *self.args) self.lazy_cube = self.cube.copy(da.asarray(self.cube.data)) + self.lazy_masked_cube = self.lazy_cube.copy(da.ma.masked_array(self.cube.data)) self.lazy_regridder = Regridder(self.lazy_cube, self.lazy_cube, *self.args) def test_lazy_regrid(self): result = self.lazy_regridder(self.lazy_cube) self.assertTrue(result.has_lazy_data()) + self.assertTrue( + isinstance(da.utils.meta_from_array(result.core_data()), np.ndarray) + ) expected = self.regridder(self.cube) self.assertTrue(result == expected) + + def test_lazy_masked_regrid(self): + result = self.lazy_regridder(self.lazy_masked_cube) + self.assertTrue(result.has_lazy_data()) self.assertTrue( isinstance(da.utils.meta_from_array(result.core_data()), np.ma.MaskedArray) ) + expected = self.regridder(self.cube) + self.assertTrue(result == expected) class Test___call____invalid_types(tests.IrisTest): diff --git a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py index c0f3a49f27..8ddc1114a6 100644 --- a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py +++ b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py @@ -32,7 +32,19 @@ def create_mock_cube(array): class Test_map_complete_blocks(tests.IrisTest): def setUp(self): self.array = np.arange(8).reshape(2, 4) - self.func = lambda chunk: chunk + 1 + + def func(chunk): + """Use a function that cannot be 'sampled'. + + To make sure the call to map_blocks is correct for any function, + we define this function that cannot be called with size 0 arrays + to infer the output meta. + """ + if chunk.size == 0: + raise ValueError + return chunk + 1 + + self.func = func self.func_result = self.array + 1 def test_non_lazy_input(self): @@ -65,6 +77,7 @@ def test_dask_masked_array_input(self): array = da.ma.masked_array(np.arange(2), mask=np.arange(2)) result = map_complete_blocks(array, self.func, dims=tuple(), out_sizes=tuple()) self.assertTrue(is_lazy_data(result)) + self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ma.MaskedArray)) self.assertArrayEqual(result.compute(), np.ma.masked_array([1, 2], mask=[0, 1])) def test_rechunk(self): From 63e0d1c44ec9fb4f2b657e97a5b740045be0da31 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 3 Jul 2024 10:56:00 +0200 Subject: [PATCH 3/5] Add more tests --- lib/iris/_lazy_data.py | 4 ++- .../lazy_data/test_map_complete_blocks.py | 34 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 1e094d8ebe..4a9f26d98b 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -601,7 +601,9 @@ def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): # Assume operation does not change dtype and meta if not specified. if "meta" not in kwargs: kwargs["meta"] = da.utils.meta_from_array(data) - if "dtype" not in kwargs: + if "dtype" in kwargs: + kwargs["meta"] = kwargs["meta"].astype(kwargs["dtype"]) + else: kwargs["dtype"] = kwargs["meta"].dtype result = data.map_blocks(func, *args, chunks=out_chunks, **kwargs) diff --git a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py index 8ddc1114a6..4fa2cf142e 100644 --- a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py +++ b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py @@ -80,6 +80,40 @@ def test_dask_masked_array_input(self): self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ma.MaskedArray)) self.assertArrayEqual(result.compute(), np.ma.masked_array([1, 2], mask=[0, 1])) + def test_dask_array_input_with_meta(self): + lazy_array = da.asarray(self.array, chunks=((1, 1), (4,))) + meta = np.empty((), dtype=np.float32) + + def func(chunk): + if chunk.size == 0: + raise ValueError + return (chunk + 1).astype(np.float32) + + result = map_complete_blocks( + lazy_array, func, dims=(1,), out_sizes=(4,), meta=meta + ) + self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ndarray)) + self.assertTrue(result.dtype == meta.dtype) + self.assertTrue(result.compute().dtype == meta.dtype) + self.assertArrayEqual(result.compute(), self.func_result) + + def test_dask_array_input_with_dtype(self): + lazy_array = da.ma.masked_array(self.array, chunks=((1, 1), (4,))) + dtype = np.float32 + + def func(chunk): + if chunk.size == 0: + raise ValueError + return (chunk + 1).astype(np.float32) + + result = map_complete_blocks( + lazy_array, func, dims=(1,), out_sizes=(4,), dtype=dtype + ) + self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ma.MaskedArray)) + self.assertTrue(result.dtype == dtype) + self.assertTrue(result.compute().dtype == dtype) + self.assertArrayEqual(result.compute(), self.func_result) + def test_rechunk(self): lazy_array = da.asarray(self.array, chunks=((1, 1), (2, 2))) cube, _ = create_mock_cube(lazy_array) From edaa84acd234273efaa92dec38426ec63a39469d Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 3 Jul 2024 10:59:29 +0200 Subject: [PATCH 4/5] Add whatsnew --- docs/src/whatsnew/latest.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index fb8098f766..bd5d2383b2 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -37,7 +37,7 @@ This document explains the changes made to Iris for this release ============= #. `@bouweandela`_ updated the ``chunktype`` of Dask arrays, so it corresponds - to the array content. (:pull:`5801`) + to the array content. (:pull:`5801`) and (:pull:`5989`) #. `@rcomer`_ made the :obj:`~iris.analysis.WPERCENTILE` aggregator work with :func:`~iris.cube.Cube.rolling_window`. (:issue:`5777`, :pull:`5825`) From c1e5e1ae6fbcf704fc92c3dea205658ae752f7b7 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 10 Jul 2024 11:34:29 +0200 Subject: [PATCH 5/5] Specify dtype in map_complete_blocks and undo change to lazy_elementwise --- lib/iris/_lazy_data.py | 25 +++++----- lib/iris/analysis/__init__.py | 7 +-- lib/iris/analysis/_area_weighted.py | 13 +++-- lib/iris/analysis/_regrid.py | 13 +++-- .../regrid/test_RectilinearRegridder.py | 12 ++--- .../lazy_data/test_map_complete_blocks.py | 47 +++++++++---------- 6 files changed, 63 insertions(+), 54 deletions(-) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index bf0af5cc29..a3dfa1edb4 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -536,15 +536,13 @@ def lazy_elementwise(lazy_array, elementwise_op): # This makes good practical sense for unit conversions, as a Unit.convert # call may cast to float, or not, depending on unit equality : Thus, it's # much safer to get udunits to decide that for us. - meta = da.utils.meta_from_array(lazy_array) - new_meta = elementwise_op(meta) + dtype = elementwise_op(np.zeros(1, lazy_array.dtype)).dtype + meta = da.utils.meta_from_array(lazy_array).astype(dtype) - return da.map_blocks( - elementwise_op, lazy_array, dtype=new_meta.dtype, meta=new_meta - ) + return da.map_blocks(elementwise_op, lazy_array, dtype=dtype, meta=meta) -def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): +def map_complete_blocks(src, func, dims, out_sizes, dtype, *args, **kwargs): """Apply a function to complete blocks. Complete means that the data is not chunked along the chosen dimensions. @@ -560,6 +558,8 @@ def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): Dimensions that cannot be chunked. out_sizes : tuple of int Output size of dimensions that cannot be chunked. + dtype : + Output dtype. *args : tuple Additional arguments to pass to `func`. **kwargs : dict @@ -599,14 +599,11 @@ def map_complete_blocks(src, func, dims, out_sizes, *args, **kwargs): for dim, size in zip(dims, out_sizes): out_chunks[dim] = size - # Assume operation does not change dtype and meta if not specified. - if "meta" not in kwargs: - kwargs["meta"] = da.utils.meta_from_array(data) - if "dtype" in kwargs: - kwargs["meta"] = kwargs["meta"].astype(kwargs["dtype"]) - else: - kwargs["dtype"] = kwargs["meta"].dtype + # Assume operation preserves mask. + meta = da.utils.meta_from_array(data).astype(dtype) - result = data.map_blocks(func, *args, chunks=out_chunks, **kwargs) + result = data.map_blocks( + func, *args, chunks=out_chunks, meta=meta, dtype=dtype, **kwargs + ) return result diff --git a/lib/iris/analysis/__init__.py b/lib/iris/analysis/__init__.py index 5d44e563c0..ca3b820cb3 100644 --- a/lib/iris/analysis/__init__.py +++ b/lib/iris/analysis/__init__.py @@ -1390,9 +1390,10 @@ def _percentile(data, percent, fast_percentile_method=False, **kwargs): result = iris._lazy_data.map_complete_blocks( data, - _calc_percentile, - (-1,), - percent.shape, + func=_calc_percentile, + dims=(-1,), + out_sizes=percent.shape, + dtype=np.float64, percent=percent, fast_percentile_method=fast_percentile_method, **kwargs, diff --git a/lib/iris/analysis/_area_weighted.py b/lib/iris/analysis/_area_weighted.py index a25a21bb47..3ed4f8aa33 100644 --- a/lib/iris/analysis/_area_weighted.py +++ b/lib/iris/analysis/_area_weighted.py @@ -392,11 +392,18 @@ def _regrid_area_weighted_rectilinear_src_and_grid__perform( tgt_shape = (len(grid_y.points), len(grid_x.points)) + # Specify the output dtype + if np.issubdtype(src_cube.dtype, np.integer): + out_dtype = np.float64 + else: + out_dtype = src_cube.dtype + new_data = map_complete_blocks( src_cube, - _regrid_along_dims, - (src_y_dim, src_x_dim), - meshgrid_x.shape, + func=_regrid_along_dims, + dims=(src_y_dim, src_x_dim), + out_sizes=meshgrid_x.shape, + dtype=out_dtype, x_dim=src_x_dim, y_dim=src_y_dim, weights=weights, diff --git a/lib/iris/analysis/_regrid.py b/lib/iris/analysis/_regrid.py index 6c10b8c404..24b88dbb86 100644 --- a/lib/iris/analysis/_regrid.py +++ b/lib/iris/analysis/_regrid.py @@ -935,11 +935,18 @@ def __call__(self, src): x_dim = src.coord_dims(src_x_coord)[0] y_dim = src.coord_dims(src_y_coord)[0] + # Specify the output dtype + if self._method == "linear" and np.issubdtype(src.dtype, np.integer): + out_dtype = np.float64 + else: + out_dtype = src.dtype + data = map_complete_blocks( src, - self._regrid, - (y_dim, x_dim), - sample_grid_x.shape, + func=self._regrid, + dims=(y_dim, x_dim), + out_sizes=sample_grid_x.shape, + dtype=out_dtype, x_dim=x_dim, y_dim=y_dim, src_x_coord=src_x_coord, diff --git a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py index 474ffb02cc..3f841b938a 100644 --- a/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py +++ b/lib/iris/tests/unit/analysis/regrid/test_RectilinearRegridder.py @@ -480,19 +480,19 @@ def setUp(self): def test_lazy_regrid(self): result = self.lazy_regridder(self.lazy_cube) self.assertTrue(result.has_lazy_data()) - self.assertTrue( - isinstance(da.utils.meta_from_array(result.core_data()), np.ndarray) - ) + meta = da.utils.meta_from_array(result.core_data()) + self.assertTrue(meta.__class__ is np.ndarray) expected = self.regridder(self.cube) + self.assertEqual(result.dtype, expected.dtype) self.assertTrue(result == expected) def test_lazy_masked_regrid(self): result = self.lazy_regridder(self.lazy_masked_cube) self.assertTrue(result.has_lazy_data()) - self.assertTrue( - isinstance(da.utils.meta_from_array(result.core_data()), np.ma.MaskedArray) - ) + meta = da.utils.meta_from_array(result.core_data()) + self.assertTrue(isinstance(meta, np.ma.MaskedArray)) expected = self.regridder(self.cube) + self.assertEqual(result.dtype, expected.dtype) self.assertTrue(result == expected) diff --git a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py index 4fa2cf142e..7d619353ed 100644 --- a/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py +++ b/lib/iris/tests/unit/lazy_data/test_map_complete_blocks.py @@ -50,7 +50,9 @@ def func(chunk): def test_non_lazy_input(self): # Check that a non-lazy input doesn't trip up the functionality. cube, cube_data = create_mock_cube(self.array) - result = map_complete_blocks(cube, self.func, dims=(1,), out_sizes=(4,)) + result = map_complete_blocks( + cube, self.func, dims=(1,), out_sizes=(4,), dtype=self.array.dtype + ) self.assertFalse(is_lazy_data(result)) self.assertArrayEqual(result, self.func_result) # check correct data was accessed @@ -60,7 +62,9 @@ def test_non_lazy_input(self): def test_lazy_input(self): lazy_array = da.asarray(self.array, chunks=((1, 1), (4,))) cube, cube_data = create_mock_cube(lazy_array) - result = map_complete_blocks(cube, self.func, dims=(1,), out_sizes=(4,)) + result = map_complete_blocks( + cube, self.func, dims=(1,), out_sizes=(4,), dtype=lazy_array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), self.func_result) # check correct data was accessed @@ -69,35 +73,22 @@ def test_lazy_input(self): def test_dask_array_input(self): lazy_array = da.asarray(self.array, chunks=((1, 1), (4,))) - result = map_complete_blocks(lazy_array, self.func, dims=(1,), out_sizes=(4,)) + result = map_complete_blocks( + lazy_array, self.func, dims=(1,), out_sizes=(4,), dtype=lazy_array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), self.func_result) def test_dask_masked_array_input(self): array = da.ma.masked_array(np.arange(2), mask=np.arange(2)) - result = map_complete_blocks(array, self.func, dims=tuple(), out_sizes=tuple()) + result = map_complete_blocks( + array, self.func, dims=tuple(), out_sizes=tuple(), dtype=array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ma.MaskedArray)) self.assertArrayEqual(result.compute(), np.ma.masked_array([1, 2], mask=[0, 1])) - def test_dask_array_input_with_meta(self): - lazy_array = da.asarray(self.array, chunks=((1, 1), (4,))) - meta = np.empty((), dtype=np.float32) - - def func(chunk): - if chunk.size == 0: - raise ValueError - return (chunk + 1).astype(np.float32) - - result = map_complete_blocks( - lazy_array, func, dims=(1,), out_sizes=(4,), meta=meta - ) - self.assertTrue(isinstance(da.utils.meta_from_array(result), np.ndarray)) - self.assertTrue(result.dtype == meta.dtype) - self.assertTrue(result.compute().dtype == meta.dtype) - self.assertArrayEqual(result.compute(), self.func_result) - - def test_dask_array_input_with_dtype(self): + def test_dask_array_input_with_different_output_dtype(self): lazy_array = da.ma.masked_array(self.array, chunks=((1, 1), (4,))) dtype = np.float32 @@ -117,7 +108,9 @@ def func(chunk): def test_rechunk(self): lazy_array = da.asarray(self.array, chunks=((1, 1), (2, 2))) cube, _ = create_mock_cube(lazy_array) - result = map_complete_blocks(cube, self.func, dims=(1,), out_sizes=(4,)) + result = map_complete_blocks( + cube, self.func, dims=(1,), out_sizes=(4,), dtype=lazy_array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), self.func_result) @@ -129,7 +122,9 @@ def func(_): return np.arange(2).reshape(1, 2) func_result = [[0, 1], [0, 1]] - result = map_complete_blocks(cube, func, dims=(1,), out_sizes=(2,)) + result = map_complete_blocks( + cube, func, dims=(1,), out_sizes=(2,), dtype=lazy_array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), func_result) @@ -137,7 +132,9 @@ def test_multidimensional_input(self): array = np.arange(2 * 3 * 4).reshape(2, 3, 4) lazy_array = da.asarray(array, chunks=((1, 1), (1, 2), (4,))) cube, _ = create_mock_cube(lazy_array) - result = map_complete_blocks(cube, self.func, dims=(1, 2), out_sizes=(3, 4)) + result = map_complete_blocks( + cube, self.func, dims=(1, 2), out_sizes=(3, 4), dtype=lazy_array.dtype + ) self.assertTrue(is_lazy_data(result)) self.assertArrayEqual(result.compute(), array + 1)