diff --git a/xrspatial/focal.py b/xrspatial/focal.py index 2cb276fb..dae8c4dd 100644 --- a/xrspatial/focal.py +++ b/xrspatial/focal.py @@ -26,7 +26,7 @@ class cupy(object): ndarray = False from xrspatial.convolution import (_available_memory_bytes, _convolve_2d_cupy, _convolve_2d_numpy, - convolve_2d, custom_kernel) + _promote_float, convolve_2d, custom_kernel) from xrspatial.dataset_support import supports_dataset from xrspatial.utils import (ArrayTypeFunctionMapping, _boundary_to_dask, _pad_array, _validate_boundary, _validate_raster, _validate_scalar, cuda_args, @@ -460,8 +460,7 @@ def _calc_variety(array): @ngjit def _apply_numpy(data, kernel, func): - data = data.astype(np.float32) - + # Caller must promote ``data`` to a float dtype (see ``_promote_float``). out = np.zeros_like(data) rows, cols = data.shape krows, kcols = kernel.shape @@ -483,6 +482,7 @@ def _apply_numpy(data, kernel, func): def _apply_numpy_boundary(data, kernel, func, boundary='nan'): + data = data.astype(_promote_float(data.dtype)) if boundary == 'nan': return _apply_numpy(data, kernel, func) pad_h = kernel.shape[0] // 2 @@ -497,7 +497,7 @@ def _apply_numpy_boundary(data, kernel, func, boundary='nan'): def _apply_dask_numpy(data, kernel, func, boundary='nan'): - data = data.astype(np.float32) + data = data.astype(_promote_float(data.dtype)) _func = partial(_apply_numpy, kernel=kernel, func=func) pad_h = kernel.shape[0] // 2 @@ -511,7 +511,7 @@ def _apply_dask_numpy(data, kernel, func, boundary='nan'): def _apply_cupy(data, kernel, func): - return _focal_stats_func_cupy(data.astype(cupy.float32), kernel, func) + return _focal_stats_func_cupy(data.astype(_promote_float(data.dtype)), kernel, func) def _apply_cupy_boundary(data, kernel, func, boundary='nan'): @@ -529,7 +529,7 @@ def _apply_cupy_boundary(data, kernel, func, boundary='nan'): def _apply_dask_cupy(data, kernel, func, boundary='nan'): - data = data.astype(cupy.float32) + data = data.astype(_promote_float(data.dtype)) pad_h = kernel.shape[0] // 2 pad_w = kernel.shape[1] // 2 _func = partial(_focal_stats_func_cupy, kernel=kernel, func=func) @@ -982,7 +982,7 @@ def _focal_sum_cuda(data, kernel, out): def _focal_stats_func_cupy(data, kernel, func=_focal_max_cuda): kernel = cupy.asarray(kernel) - out = cupy.empty(data.shape, dtype='f4') + out = cupy.empty(data.shape, dtype=_promote_float(data.dtype)) out[:, :] = cupy.nan griddim, blockdim = cuda_args(data.shape) func[griddim, blockdim](data, kernel, cupy.asarray(out)) @@ -1038,7 +1038,7 @@ def _focal_stats_cupy(agg, kernel, stats_funcs): ) stats_aggs = [] for stats in stats_funcs: - data = agg.data.astype(cupy.float32) + data = agg.data.astype(_promote_float(agg.data.dtype)) stats_data = _stats_cupy_mapper[stats](data, kernel) stats_agg = xr.DataArray( stats_data, @@ -1089,7 +1089,7 @@ def _focal_stats_dask_cupy(agg, kernel, stats_funcs, boundary='nan'): for stat_name in stats_funcs: cuda_kernel = _stats_cuda_mapper[stat_name] _func = partial(_focal_stats_func_cupy, kernel=kernel, func=cuda_kernel) - data = agg.data.astype(cupy.float32) + data = agg.data.astype(_promote_float(agg.data.dtype)) stats_data = data.map_overlap( _func, depth=(pad_h, pad_w), boundary=dask_bnd, meta=cupy.array(())) diff --git a/xrspatial/tests/test_focal.py b/xrspatial/tests/test_focal.py index fdb7c67e..8e90158b 100644 --- a/xrspatial/tests/test_focal.py +++ b/xrspatial/tests/test_focal.py @@ -558,6 +558,93 @@ def test_focal_stats_dask_cupy(): equal_nan=True, rtol=1e-4) +# --- float64 preservation (issue-2769) ------------------------------------ +# apply() and focal_stats() used to cast every input to float32 internally, +# silently downcasting float64 rasters. convolve_2d() preserves the input +# floating dtype (see test_convolution); these mirror that contract for the +# focal APIs across all four backends. + + +def _compute_dtype(agg): + """Return the materialised dtype regardless of backend.""" + data = agg.data + if hasattr(data, 'compute'): + data = data.compute() + if hasattr(data, 'get'): + data = data.get() + return data.dtype + + +@pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy']) +def test_apply_preserves_float64(backend): + from xrspatial.tests.general_checks import has_cuda_and_cupy + if 'cupy' in backend and not has_cuda_and_cupy(): + pytest.skip("Requires CUDA and CuPy") + if 'dask' in backend and da is None: + pytest.skip("Requires Dask") + + data = np.arange(20, dtype=np.float64).reshape(4, 5) + kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]) + + if 'cupy' in backend: + from xrspatial.focal import _focal_mean_cuda + func = _focal_mean_cuda + else: + from xrspatial.focal import _calc_mean + func = _calc_mean + + agg = create_test_raster(data, backend=backend, chunks=(2, 3)) + result = apply(agg, kernel, func) + + assert _compute_dtype(result) == np.float64 + + +@pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy']) +def test_focal_stats_preserves_float64(backend): + from xrspatial.tests.general_checks import has_cuda_and_cupy + if 'cupy' in backend and not has_cuda_and_cupy(): + pytest.skip("Requires CUDA and CuPy") + if 'dask' in backend and da is None: + pytest.skip("Requires Dask") + + data = np.arange(20, dtype=np.float64).reshape(4, 5) + kernel = custom_kernel(np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]])) + + agg = create_test_raster(data, backend=backend, chunks=(2, 3)) + result = focal_stats(agg, kernel, + stats_funcs=['mean', 'sum', 'min', 'max', + 'std', 'var', 'range']) + + assert _compute_dtype(result) == np.float64 + + +@pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy']) +def test_apply_keeps_float32(backend): + # The other side of the contract: a float32 input must not be promoted + # to float64. (On dask the lazy dtype is float64, but the computed + # result is float32 -- matching convolve_2d.) + from xrspatial.tests.general_checks import has_cuda_and_cupy + if 'cupy' in backend and not has_cuda_and_cupy(): + pytest.skip("Requires CUDA and CuPy") + if 'dask' in backend and da is None: + pytest.skip("Requires Dask") + + data = np.arange(20, dtype=np.float32).reshape(4, 5) + kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]) + + if 'cupy' in backend: + from xrspatial.focal import _focal_mean_cuda + func = _focal_mean_cuda + else: + from xrspatial.focal import _calc_mean + func = _calc_mean + + agg = create_test_raster(data, backend=backend, chunks=(2, 3)) + result = apply(agg, kernel, func) + + assert _compute_dtype(result) == np.float32 + + # --- focal_stats NaN handling (issue-1092) -------------------------------- @pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy'])