validate_arrays: ensure chunksizes of arrays are matching (#577)

Author: thuydotm

xrspatial/classify.py

@@ -585,7 +585,7 @@ def _run_cupy_jenks_matrices(data, n_classes):
     nl = data.shape[0] + 1
     variance = 0.0
 
-    for l in range(2, nl): # noqa
+    for l in range(2, nl):  # noqa
         sum = 0.0
         sum_squares = 0.0
         w = 0.0

xrspatial/utils.py

@@ -129,6 +129,12 @@ def validate_arrays(*arrays):
         if not isinstance(first_array.data, type(arrays[i].data)):
             raise ValueError("input arrays must have same type")
 
+    # ensure dask chunksizes of all arrays are the same
+    if isinstance(first_array.data, da.Array):
+        for i in range(1, len(arrays)):
+            if first_array.chunks != arrays[i].chunks:
+                arrays[i].data = arrays[i].data.rechunk(first_array.chunks)
+
 
 def get_xy_range(raster, xdim=None, ydim=None):
     """

xrspatial/zonal.py

@@ -10,7 +10,7 @@
 import dask.dataframe as dd
 from dask import delayed
 
-from xrspatial.utils import ngjit
+from xrspatial.utils import ngjit, validate_arrays
 
 
 def _stats_count(data):
@@ -472,14 +472,13 @@ def stats(
         3    30        3850
     """
 
-    if zones.shape != values.shape:
-        raise ValueError("`zones` and `values` must have same shape.")
+    validate_arrays(zones, values)
 
     if not (
         issubclass(zones.data.dtype.type, np.integer)
         or issubclass(zones.data.dtype.type, np.floating)
     ):
-        raise ValueError("`zones` must be an array of integers.")
+        raise ValueError("`zones` must be an array of integers or floats.")
 
     if not (
         issubclass(values.data.dtype.type, np.integer)
@@ -521,11 +520,6 @@ def stats(
         )
     else:
         # dask case
-
-        # make sure chunksizes of `zones` and `values` are matching
-        if zones.chunks != values.chunks:
-            values.data = values.data.rechunk(zones.chunks)
-
         stats_df = _stats_dask_numpy(
             zones.data,
             values.data,