WIP cupy support for engine="flox"

Author: dcherian

flox/aggregate_flox.py

@@ -4,6 +4,18 @@
 
 from .xrutils import isnull
 
+import cupy as cp
+import cupyx
+
+cupy_array_type = (cp.ndarray,)
+
+cupy_ops = {
+    np.add: cupyx.scatter_add,
+    np.minimum: cupyx.scatter_min,
+    np.maximum: cupyx.scatter_max,
+}
+
+
 def _prepare_for_flox(group_idx, array):
     """
     Sort the input array once to save time.
@@ -24,7 +36,9 @@ def _np_grouped_op(group_idx, array, op, axis=-1, size=None, fill_value=None, dt
     most of this code is from shoyer's gist
     https://gist.github.com/shoyer/f538ac78ae904c936844
     """
-    # assumes input is sorted, which I do in core._prepare_for_flox
+    # For numpy arrays, assumes input is sorted, which I do in _prepare_for_flox
+    # For cupy arrays, sorting is not needed
+
     aux = group_idx
 
     flag = np.concatenate((np.array([True], like=array), aux[1:] != aux[:-1]))
@@ -37,7 +51,12 @@ def _np_grouped_op(group_idx, array, op, axis=-1, size=None, fill_value=None, dt
         dtype = array.dtype
 
     if out is None:
-        out = np.full(array.shape[:-1] + (size,), fill_value=fill_value, dtype=dtype)
+        out = np.full(array.shape[:-1] + (size,), fill_value=fill_value, dtype=dtype, like=array)
+
+    if isinstance(array, cupy_array_type):
+        op = cupy_ops[op]
+        op(out, group_idx, array)
+        return out
 
     if (len(uniques) == size) and (uniques == np.arange(size, like=array)).all():
         # The previous version of this if condition