Use npg sum_of_squares

asv_bench/benchmarks/reduce.py

@@ -6,7 +6,7 @@
 from . import parameterized
 
 N = 1000
-funcs = ["sum", "nansum", "mean", "nanmean", "max"]
+funcs = ["sum", "nansum", "mean", "nanmean", "max", "var", "nanvar"]
 engines = ["flox", "numpy"]
 
 

ci/docs.yml

@@ -7,7 +7,7 @@ dependencies:
   - xarray
   - numpy>=1.20
   - numpydoc
-  - numpy_groupies
+  - numpy_groupies>=0.9.19
   - toolz
   - matplotlib-base
   - myst-parser

ci/environment.yml

@@ -16,7 +16,7 @@ dependencies:
   - pytest-xdist
   - xarray
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba

ci/minimal-requirements.yml

@@ -9,7 +9,7 @@ dependencies:
   - pytest-cov
   - pytest-xdist
   - numpy==1.20
-  - numpy_groupies==0.9.15
+  - numpy_groupies==0.9.19
   - pandas
   - pooch
   - toolz

ci/no-dask.yml

@@ -13,7 +13,7 @@ dependencies:
   - xarray
   - numpydoc
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba

ci/no-xarray.yml

@@ -13,7 +13,7 @@ dependencies:
   - dask-core
   - numpydoc
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba

flox/aggregate_npg.py

@@ -9,14 +9,43 @@ def _get_aggregate(engine):
 
 
 def sum_of_squares(
-    group_idx, array, engine, *, axis=-1, func="sum", size=None, fill_value=None, dtype=None
+    group_idx,
+    array,
+    engine,
+    *,
+    axis=-1,
+    size=None,
+    fill_value=None,
+    dtype=None,
 ):
 
     return _get_aggregate(engine).aggregate(
         group_idx,
-        array**2,
+        array,
         axis=axis,
-        func=func,
+        func="sumofsquares",
+        size=size,
+        fill_value=fill_value,
+        dtype=dtype,
+    )
+
+
+def nansum_of_squares(
+    group_idx,
+    array,
+    engine,
+    *,
+    axis=-1,
+    size=None,
+    fill_value=None,
+    dtype=None,
+):
+
+    return _get_aggregate(engine).aggregate(
+        group_idx,
+        array,
+        axis=axis,
+        func="nansumofsquares",
         size=size,
         fill_value=fill_value,
         dtype=dtype,
@@ -55,19 +84,6 @@ def nanprod(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dt
     )
 
 
-def nansum_of_squares(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dtype=None):
-    return sum_of_squares(
-        group_idx,
-        array,
-        engine=engine,
-        func="nansum",
-        size=size,
-        fill_value=fill_value,
-        axis=axis,
-        dtype=dtype,
-    )
-
-
 def _len(group_idx, array, engine, *, func, axis=-1, size=None, fill_value=None, dtype=None):
     result = _get_aggregate(engine).aggregate(
         group_idx,

flox/aggregations.py

@@ -39,7 +39,7 @@ def generic_aggregate(
             method_ = getattr(aggregate_npg, func)
             method = partial(method_, engine=engine)
         except AttributeError:
-            aggregate = npg.aggregate_np if engine == "numpy" else npg.aggregate_nb
+            aggregate = aggregate_npg._get_aggregate(engine).aggregate
             method = partial(aggregate, func=func)
     else:
         raise ValueError(

setup.cfg

@@ -28,7 +28,7 @@ python_requires = >=3.8
 install_requires =
     pandas
     numpy >= '1.20'
-    numpy_groupies >= '0.9.15'
+    numpy_groupies >= '0.9.19'
     toolz
 
 [options.extras_require]