Merge branch 'main' into numbagg

* main:
  Fix benchmarks
  Fix engine='numba' (#73)
  Switch to pre-commit.ci (#71)

Author: dcherian

.github/workflows/linting.yaml

@@ -13,7 +13,7 @@ def setup(self, *args, **kwargs):
 
     @parameterized("kind", ("cohorts", "mapreduce"))
     def time_combine(self, kind):
-        flox.core._npg_combine(
+        flox.core._grouped_combine(
             getattr(self, f"x_chunk_{kind}"),
             **self.kwargs,
             keepdims=True,
@@ -22,7 +22,7 @@ def time_combine(self, kind):
 
     @parameterized("kind", ("cohorts", "mapreduce"))
     def peakmem_combine(self, kind):
-        flox.core._npg_combine(
+        flox.core._grouped_combine(
             getattr(self, f"x_chunk_{kind}"),
             **self.kwargs,
             keepdims=True,
@@ -58,4 +58,4 @@ def construct_member(groups):
         ]
 
         self.x_chunk_cohorts = [construct_member(groups) for groups in [np.array((1, 2, 3, 4))] * 4]
-        self.kwargs = {"agg": flox.aggregations.mean, "axis": (3,), "group_ndim": 1}
+        self.kwargs = {"agg": flox.aggregations.mean, "axis": (3,), "neg_axis": (-1,)}

benchmarks/combine.py

@@ -2,11 +2,15 @@
 import numpy_groupies as npg
 
 
+def _get_aggregate(engine):
+    return npg.aggregate_numpy if engine == "numpy" else npg.aggregate_numba
+
+
 def sum_of_squares(
-    group_idx, array, *, axis=-1, func="sum", size=None, fill_value=None, dtype=None
+    group_idx, array, engine, *, axis=-1, func="sum", size=None, fill_value=None, dtype=None
 ):
 
-    return npg.aggregate_numpy.aggregate(
+    return _get_aggregate(engine).aggregate(
         group_idx,
         array**2,
         axis=axis,
@@ -17,12 +21,12 @@ def sum_of_squares(
     )
 
 
-def nansum(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
+def nansum(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dtype=None):
     # npg takes out NaNs before calling np.bincount
     # This means that all NaN groups are equivalent to absent groups
     # This behaviour does not work for xarray
 
-    return npg.aggregate_numpy.aggregate(
+    return _get_aggregate(engine).aggregate(
         group_idx,
         np.where(np.isnan(array), 0, array),
         axis=axis,
@@ -33,12 +37,12 @@ def nansum(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None)
     )
 
 
-def nanprod(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
+def nanprod(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dtype=None):
     # npg takes out NaNs before calling np.bincount
     # This means that all NaN groups are equivalent to absent groups
     # This behaviour does not work for xarray
 
-    return npg.aggregate_numpy.aggregate(
+    return _get_aggregate(engine).aggregate(
         group_idx,
         np.where(np.isnan(array), 1, array),
         axis=axis,
@@ -49,7 +53,14 @@ def nanprod(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None
     )
 
 
-def nansum_of_squares(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
+def nansum_of_squares(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dtype=None):
     return sum_of_squares(
-        group_idx, array, func="nansum", size=size, fill_value=fill_value, axis=axis, dtype=dtype
+        group_idx,
+        array,
+        engine=engine,
+        func="nansum",
+        size=size,
+        fill_value=fill_value,
+        axis=axis,
+        dtype=dtype,
     )

flox/aggregate_npg.py

@@ -10,31 +10,38 @@
 
 
 def generic_aggregate(
-    group_idx, array, *, engine, func, axis=-1, size=None, fill_value=None, dtype=None, **kwargs
+    group_idx,
+    array,
+    *,
+    engine: str,
+    func: str,
+    axis=-1,
+    size=None,
+    fill_value=None,
+    dtype=None,
+    **kwargs,
 ):
     if engine == "flox":
         try:
             method = getattr(aggregate_flox, func)
         except AttributeError:
             method = partial(npg.aggregate_numpy.aggregate, func=func)
+
     elif engine == "numbagg":
         from . import aggregate_numbagg
 
         try:
             method = getattr(aggregate_numbagg, func)
         except AttributeError:
             method = partial(npg.aggregate_numpy.aggregate, func=func)
-    elif engine == "numpy":
-        try:
-            # TODO: fix numba here
-            method = getattr(aggregate_npg, func)
-        except AttributeError:
-            method = partial(npg.aggregate_np, func=func)
-    elif engine == "numba":
+
+    elif engine in ["numpy", "numba"]:
         try:
-            method = getattr(aggregate_npg, f"{func}")
+            method_ = getattr(aggregate_npg, func)
+            method = partial(method_, engine=engine)
         except AttributeError:
-            method = partial(npg.aggregate_nb, func=func)
+            aggregate = npg.aggregate_np if engine == "numpy" else npg.aggregate_nb
+            method = partial(aggregate, func=func)
     else:
         raise ValueError(
             f"Expected engine to be one of ['flox', 'numpy', 'numba', 'numbagg']. Received {engine} instead."