Name dask tasks for remaining xrspatial tool modules (#3256) (#3326)

Author: brendancol

xrspatial/classify.py

@@ -26,6 +26,7 @@ class cupy(object):
 
 from xrspatial.utils import (
     ArrayTypeFunctionMapping,
+    _dask_task_name_kwargs,
     _validate_raster,
     _validate_scalar,
     cuda_args,
@@ -74,7 +75,7 @@ def _run_numpy_binary(data, values):
 
 def _run_dask_numpy_binary(data, values):
     _func = partial(_run_numpy_binary, values=values)
-    out = data.map_blocks(_func)
+    out = data.map_blocks(_func, **_dask_task_name_kwargs('xrspatial.binary'))
     return out
 
 
@@ -104,7 +105,8 @@ def _run_cupy_binary(data, values):
 
 
 def _run_dask_cupy_binary(data, values_cupy):
-    out = data.map_blocks(lambda da: _run_cupy_binary(da, values_cupy), meta=cupy.array(()))
+    out = data.map_blocks(lambda da: _run_cupy_binary(da, values_cupy), meta=cupy.array(()),
+                          **_dask_task_name_kwargs('xrspatial.binary'))
     return out
 
 
@@ -224,7 +226,7 @@ def _run_dask_numpy_bin(data, bins, new_values):
                     bins=bins,
                     new_values=new_values)
 
-    out = data.map_blocks(_func)
+    out = data.map_blocks(_func, **_dask_task_name_kwargs('xrspatial.reclassify'))
     return out
 
 
@@ -281,7 +283,8 @@ def _run_cupy_bin(data, bins, new_values):
 def _run_dask_cupy_bin(data, bins_cupy, new_values_cupy):
     out = data.map_blocks(lambda da:
                           _run_cupy_bin(da, bins_cupy, new_values_cupy),
-                          meta=cupy.array(()))
+                          meta=cupy.array(()),
+                          **_dask_task_name_kwargs('xrspatial.reclassify'))
     return out
 
 

xrspatial/cost_distance.py

@@ -49,6 +49,7 @@ class cupy:  # type: ignore[no-redef]
         ndarray = False
 
 from xrspatial.utils import (
+    _dask_task_name_kwargs,
     _validate_raster,
     cuda_args, get_dataarray_resolution, ngjit,
     has_cuda_and_cupy, is_cupy_array, is_dask_cupy,
@@ -506,6 +507,7 @@ def _chunk_func(source_block, friction_block):
             boundary=np.nan,
             dtype=np.float32,
             meta=cp.array((), dtype=cp.float32),
+            **_dask_task_name_kwargs('xrspatial.cost_distance'),
         )
 
     # Unbounded or padding too large: convert to dask+numpy, use CPU path
@@ -1187,6 +1189,7 @@ def _cost_distance_dask(source_da, friction_da, cellsize_x, cellsize_y,
         boundary=np.nan,
         dtype=np.float32,
         meta=np.array((), dtype=np.float32),
+        **_dask_task_name_kwargs('xrspatial.cost_distance'),
     )
     return out
 

xrspatial/dasymetric.py

@@ -52,6 +52,7 @@ class cupy:
 
 from xrspatial.utils import (
     ArrayTypeFunctionMapping,
+    _dask_task_name_kwargs,
     has_cuda_and_cupy,
     has_dask_array,
     is_cupy_array,
@@ -293,6 +294,7 @@ def _disaggregate_dask_numpy(zones_da, weight_da, values_dict, method,
         method=method,
         nodata_zone=nodata_zone,
         dtype=np.float64,
+        **_dask_task_name_kwargs('xrspatial.disaggregate'),
     )
     return result
 

xrspatial/diffusion.py

@@ -33,6 +33,7 @@ class cupy:
 from xrspatial.utils import (
     ArrayTypeFunctionMapping,
     _boundary_to_dask,
+    _dask_task_name_kwargs,
     _pad_array,
     _validate_boundary,
     _validate_raster,
@@ -267,6 +268,7 @@ def _diffuse_dask_numpy(data, alpha, steps, dt_over_dx2, boundary):
                 depth=(1, 1),
                 boundary=_boundary_to_dask(boundary),
                 meta=np.array(()),
+                **_dask_task_name_kwargs('xrspatial.diffuse'),
             )
         else:
             # Pass alpha as a second dask argument to map_overlap
@@ -278,6 +280,7 @@ def _diffuse_dask_numpy(data, alpha, steps, dt_over_dx2, boundary):
                 meta=np.array(()),
                 steps=1,
                 dt_over_dx2=dt_over_dx2,
+                **_dask_task_name_kwargs('xrspatial.diffuse'),
             )
     return u
 
@@ -318,6 +321,7 @@ def _diffuse_dask_cupy(data, alpha, steps, dt_over_dx2, boundary):
                 depth=(1, 1),
                 boundary=_boundary_to_dask(boundary, is_cupy=True),
                 meta=cp.array(()),
+                **_dask_task_name_kwargs('xrspatial.diffuse'),
             )
         else:
             u = da.map_overlap(
@@ -327,6 +331,7 @@ def _diffuse_dask_cupy(data, alpha, steps, dt_over_dx2, boundary):
                 boundary=_boundary_to_dask(boundary, is_cupy=True),
                 meta=cp.array(()),
                 dt_over_dx2=dt_over_dx2,
+                **_dask_task_name_kwargs('xrspatial.diffuse'),
             )
     return u
 

xrspatial/emerging_hotspots.py

@@ -30,6 +30,7 @@
 )
 from xrspatial.utils import (
     ArrayTypeFunctionMapping,
+    _dask_task_name_kwargs,
     _validate_boundary,
     ngjit,
     not_implemented_func,
@@ -608,6 +609,7 @@ def _emerging_hotspots_dask_numpy(raster, kernel, boundary='nan'):
         gi_zscore,
         dtype=np.int8,
         meta=np.array((), dtype=np.int8),
+        **_dask_task_name_kwargs('xrspatial.emerging_hotspots.gi_bin'),
     )
 
     # Pass 3: Mann-Kendall + classification via map_blocks
@@ -620,6 +622,7 @@ def _emerging_hotspots_dask_numpy(raster, kernel, boundary='nan'):
         drop_axis=0,
         new_axis=0,
         meta=np.array((), dtype=np.float32),
+        **_dask_task_name_kwargs('xrspatial.emerging_hotspots.mk_classify'),
     )
 
     return gi_zscore, gi_bin, mk_result
@@ -690,6 +693,7 @@ def _emerging_hotspots_dask_cupy(raster, kernel, boundary='nan'):
         gi_zscore,
         dtype=np.int8,
         meta=cupy.array((), dtype=cupy.int8),
+        **_dask_task_name_kwargs('xrspatial.emerging_hotspots.gi_bin'),
     )
 
     # Pass 3: Mann-Kendall + classification via map_blocks
@@ -701,6 +705,7 @@ def _emerging_hotspots_dask_cupy(raster, kernel, boundary='nan'):
         drop_axis=0,
         new_axis=0,
         meta=cupy.array((), dtype=cupy.float32),
+        **_dask_task_name_kwargs('xrspatial.emerging_hotspots.mk_classify'),
     )
 
     return gi_zscore, gi_bin, mk_result

xrspatial/fire.py

@@ -26,6 +26,7 @@
 from xrspatial.dataset_support import supports_dataset
 from xrspatial.utils import (
     ArrayTypeFunctionMapping,
+    _dask_task_name_kwargs,
     _validate_raster,
     _validate_scalar,
     cuda_args,
@@ -165,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, meta=np.array(()),
+                         **_dask_task_name_kwargs('xrspatial.dnbr'))
 
 
 def _dnbr_cupy(pre_data, post_data):
@@ -178,7 +180,8 @@ def _dnbr_cupy(pre_data, post_data):
 
 def _dnbr_dask_cupy(pre_data, post_data):
     return da.map_blocks(_dnbr_cupy, pre_data, post_data,
-                         dtype=cupy.float32, meta=cupy.array(()))
+                         dtype=cupy.float32, meta=cupy.array(()),
+                         **_dask_task_name_kwargs('xrspatial.dnbr'))
 
 
 def dnbr(pre_nbr_agg: xr.DataArray,
@@ -256,7 +259,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, meta=np.array(()),
+                         **_dask_task_name_kwargs('xrspatial.rdnbr'))
 
 
 def _rdnbr_cupy(dnbr_data, pre_data):
@@ -269,7 +273,8 @@ def _rdnbr_cupy(dnbr_data, pre_data):
 
 def _rdnbr_dask_cupy(dnbr_data, pre_data):
     return da.map_blocks(_rdnbr_cupy, dnbr_data, pre_data,
-                         dtype=cupy.float32, meta=cupy.array(()))
+                         dtype=cupy.float32, meta=cupy.array(()),
+                         **_dask_task_name_kwargs('xrspatial.rdnbr'))
 
 
 def rdnbr(dnbr_agg: xr.DataArray,
@@ -382,7 +387,8 @@ def _bsc_gpu(data, out):
 
 
 def _bsc_dask(data):
-    return da.map_blocks(_bsc_cpu, data, dtype=np.int8, meta=np.array(()))
+    return da.map_blocks(_bsc_cpu, data, dtype=np.int8, meta=np.array(()),
+                         **_dask_task_name_kwargs('xrspatial.burn_severity_class'))
 
 
 def _bsc_cupy(data):
@@ -394,7 +400,8 @@ def _bsc_cupy(data):
 
 def _bsc_dask_cupy(data):
     return da.map_blocks(_bsc_cupy, data,
-                         dtype=np.int8, meta=cupy.array(()))
+                         dtype=np.int8, meta=cupy.array(()),
+                         **_dask_task_name_kwargs('xrspatial.burn_severity_class'))
 
 
 @supports_dataset
@@ -464,7 +471,8 @@ 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(()))
+                         meta=np.array(()),
+                         **_dask_task_name_kwargs('xrspatial.fireline_intensity'))
 
 
 def _fli_cupy(fuel_data, spread_data, heat_content):
@@ -477,7 +485,8 @@ def _fli_cupy(fuel_data, spread_data, heat_content):
 
 def _fli_dask_cupy(fuel_data, spread_data, heat_content):
     return da.map_blocks(_fli_cupy, fuel_data, spread_data, heat_content,
-                         dtype=cupy.float32, meta=cupy.array(()))
+                         dtype=cupy.float32, meta=cupy.array(()),
+                         **_dask_task_name_kwargs('xrspatial.fireline_intensity'))
 
 
 def fireline_intensity(fuel_consumed_agg: xr.DataArray,
@@ -562,7 +571,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, meta=np.array(()),
+                         **_dask_task_name_kwargs('xrspatial.flame_length'))
 
 
 def _fl_cupy(intensity_data):
@@ -575,7 +585,8 @@ def _fl_cupy(intensity_data):
 
 def _fl_dask_cupy(intensity_data):
     return da.map_blocks(_fl_cupy, intensity_data,
-                         dtype=cupy.float32, meta=cupy.array(()))
+                         dtype=cupy.float32, meta=cupy.array(()),
+                         **_dask_task_name_kwargs('xrspatial.flame_length'))
 
 
 @supports_dataset
@@ -733,6 +744,7 @@ def _ros_dask(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(()),
+        **_dask_task_name_kwargs('xrspatial.rate_of_spread'),
     )
 
 
@@ -758,6 +770,7 @@ def _ros_dask_cupy(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,
         dtype=cupy.float32, meta=cupy.array(()),
+        **_dask_task_name_kwargs('xrspatial.rate_of_spread'),
     )
 
 
@@ -907,6 +920,7 @@ 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(()),
+        **_dask_task_name_kwargs('xrspatial.kbdi'),
     )
 
 
@@ -926,6 +940,7 @@ def _kbdi_dask_cupy(kbdi_prev_data, max_temp_data, precip_data,
         _kbdi_cupy, kbdi_prev_data, max_temp_data, precip_data,
         annual_precip,
         dtype=cupy.float32, meta=cupy.array(()),
+        **_dask_task_name_kwargs('xrspatial.kbdi'),
     )
 
 

xrspatial/flood.py

@@ -29,6 +29,7 @@
     da = None
 
 from xrspatial.utils import (
+    _dask_task_name_kwargs,
     _validate_raster,
     has_cuda_and_cupy,
     is_cupy_array,
@@ -627,6 +628,7 @@ def _apply_lut(block, lut=lut):
     return _da.map_blocks(
         _apply_lut, data, dtype=np.float64,
         meta=np.array((), dtype=np.float64),
+        **_dask_task_name_kwargs('xrspatial.vegetation_roughness_nlcd'),
     )
 
 
@@ -669,6 +671,7 @@ def _apply(block):
     return _da.map_blocks(
         _apply, data, dtype=np.float64,
         meta=np.array((), dtype=np.float64),
+        **_dask_task_name_kwargs('xrspatial.vegetation_roughness_ndvi'),
     )
 
 
@@ -795,6 +798,7 @@ def _apply(lc_block, sg_block, lut=lut):
     return _da.map_blocks(
         _apply, lc, sg, dtype=np.float64,
         meta=np.array((), dtype=np.float64),
+        **_dask_task_name_kwargs('xrspatial.vegetation_curve_number'),
     )