fire: declare float32 dtype on dask+numpy map_blocks (#3394)

The six dask+numpy map_blocks wrappers (_dnbr_dask, _rdnbr_dask,
_fli_dask, _fl_dask, _ros_dask, _kbdi_dask) passed meta=np.array(())
with no dtype= argument, so the lazy dask DataArray reported float64
while the @ngjit kernels return float32. Pass dtype=np.float32 to match
the kernels and the other three backends. _bsc_dask already set
dtype=np.int8.

Adds a parametrized regression test asserting the dask+numpy declared
dtype matches the numpy backend for every fire function.

Author: brendancol

xrspatial/fire.py

@@ -166,7 +166,8 @@ def _dnbr_gpu(pre_data, post_data, out):
 
 
 def _dnbr_dask(pre_data, post_data):
-    return da.map_blocks(_dnbr_cpu, pre_data, post_data, meta=np.array(()),
+    return da.map_blocks(_dnbr_cpu, pre_data, post_data,
+                         dtype=np.float32, meta=np.array(()),
                          **_dask_task_name_kwargs('xrspatial.dnbr'))
 
 
@@ -259,7 +260,8 @@ def _rdnbr_gpu(dnbr_data, pre_data, out):
 
 
 def _rdnbr_dask(dnbr_data, pre_data):
-    return da.map_blocks(_rdnbr_cpu, dnbr_data, pre_data, meta=np.array(()),
+    return da.map_blocks(_rdnbr_cpu, dnbr_data, pre_data,
+                         dtype=np.float32, meta=np.array(()),
                          **_dask_task_name_kwargs('xrspatial.rdnbr'))
 
 
@@ -471,7 +473,7 @@ def _fli_gpu(fuel_data, spread_data, heat_content, out):
 
 def _fli_dask(fuel_data, spread_data, heat_content):
     return da.map_blocks(_fli_cpu, fuel_data, spread_data, heat_content,
-                         meta=np.array(()),
+                         dtype=np.float32, meta=np.array(()),
                          **_dask_task_name_kwargs('xrspatial.fireline_intensity'))
 
 
@@ -571,7 +573,8 @@ def _fl_gpu(intensity_data, out):
 
 
 def _fl_dask(intensity_data):
-    return da.map_blocks(_fl_cpu, intensity_data, meta=np.array(()),
+    return da.map_blocks(_fl_cpu, intensity_data,
+                         dtype=np.float32, meta=np.array(()),
                          **_dask_task_name_kwargs('xrspatial.flame_length'))
 
 
@@ -743,7 +746,7 @@ def _ros_dask(slope_data, wind_data, moisture_data,
         _ros_cpu, slope_data, wind_data, moisture_data,
         w_0, h, M_x, beta, rho_b, Gamma, eta_s, xi, epsilon,
         C_w, B_w, E_w,
-        meta=np.array(()),
+        dtype=np.float32, meta=np.array(()),
         **_dask_task_name_kwargs('xrspatial.rate_of_spread'),
     )
 
@@ -919,7 +922,7 @@ def _kbdi_gpu(kbdi_prev_data, max_temp_data, precip_data, annual_precip, out):
 def _kbdi_dask(kbdi_prev_data, max_temp_data, precip_data, annual_precip):
     return da.map_blocks(
         _kbdi_cpu, kbdi_prev_data, max_temp_data, precip_data, annual_precip,
-        meta=np.array(()),
+        dtype=np.float32, meta=np.array(()),
         **_dask_task_name_kwargs('xrspatial.kbdi'),
     )
 

xrspatial/tests/test_fire.py

@@ -664,3 +664,66 @@ def test_dataset_flame_length_accessor(self):
         })
         result = ds.xrs.flame_length()
         assert isinstance(result, xr.Dataset)
+
+
+# ---------------------------------------------------------------------------
+# Backend dtype parity (regression for #3394)
+# ---------------------------------------------------------------------------
+# Each entry builds the function's call for a given backend so the dask path's
+# declared dtype can be compared against the numpy path. The dask map_blocks
+# wrappers used to omit dtype=, defaulting the declared dtype to float64 while
+# the kernels actually return float32.
+def _dnbr_call(backend):
+    a = create_test_raster(np.array([[0.5, 0.6]], dtype='f4'), backend)
+    b = create_test_raster(np.array([[0.1, 0.3]], dtype='f4'), backend)
+    return dnbr(a, b)
+
+
+def _rdnbr_call(backend):
+    a = create_test_raster(np.array([[0.4, 0.3]], dtype='f4'), backend)
+    b = create_test_raster(np.array([[500.0, 200.0]], dtype='f4'), backend)
+    return rdnbr(a, b)
+
+
+def _bsc_call(backend):
+    a = create_test_raster(np.array([[0.5, -0.3]], dtype='f4'), backend)
+    return burn_severity_class(a)
+
+
+def _fli_call(backend):
+    a = create_test_raster(np.array([[2.0, 0.5]], dtype='f4'), backend)
+    b = create_test_raster(np.array([[0.1, 0.2]], dtype='f4'), backend)
+    return fireline_intensity(a, b)
+
+
+def _fl_call(backend):
+    a = create_test_raster(np.array([[100.0, 500.0]], dtype='f4'), backend)
+    return flame_length(a)
+
+
+def _ros_call(backend):
+    slope = create_test_raster(np.full((2, 2), 10.0, dtype='f4'), backend)
+    wind = create_test_raster(np.full((2, 2), 10.0, dtype='f4'), backend)
+    moist = create_test_raster(np.full((2, 2), 0.06, dtype='f4'), backend)
+    return rate_of_spread(slope, wind, moist)
+
+
+def _kbdi_call(backend):
+    prev = create_test_raster(np.full((2, 2), 100.0, dtype='f4'), backend)
+    temp = create_test_raster(np.full((2, 2), 30.0, dtype='f4'), backend)
+    precip = create_test_raster(np.zeros((2, 2), dtype='f4'), backend)
+    return kbdi(prev, temp, precip, annual_precip=1500.0)
+
+
+@dask_array_available
+@pytest.mark.parametrize("call", [
+    _dnbr_call, _rdnbr_call, _bsc_call, _fli_call, _fl_call, _ros_call,
+    _kbdi_call,
+])
+def test_dask_numpy_dtype_matches_numpy(call):
+    np_r = call('numpy')
+    dk_r = call('dask+numpy')
+    # declared dtype on the lazy dask array must match numpy backend
+    assert dk_r.dtype == np_r.dtype
+    # and must match what the graph actually computes
+    assert dk_r.data.compute().dtype == np_r.dtype