Preserve input float dtype in apply() and focal_stats() (#2769)

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(()))

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'])