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)