diff --git a/asv_bench/benchmarks/reduce.py b/asv_bench/benchmarks/reduce.py index 3f4a9d02f..b53a007da 100644 --- a/asv_bench/benchmarks/reduce.py +++ b/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"] diff --git a/ci/docs.yml b/ci/docs.yml index 9cdfb38e5..be9129c09 100644 --- a/ci/docs.yml +++ b/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 diff --git a/ci/environment.yml b/ci/environment.yml index d83bd3b9c..ec7eada65 100644 --- a/ci/environment.yml +++ b/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 diff --git a/ci/minimal-requirements.yml b/ci/minimal-requirements.yml index 882c8d1fb..f23032f7f 100644 --- a/ci/minimal-requirements.yml +++ b/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 diff --git a/ci/no-dask.yml b/ci/no-dask.yml index 31ce0ade3..817be4627 100644 --- a/ci/no-dask.yml +++ b/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 diff --git a/ci/no-xarray.yml b/ci/no-xarray.yml index 25c777fa1..9f81704ee 100644 --- a/ci/no-xarray.yml +++ b/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 diff --git a/flox/aggregate_npg.py b/flox/aggregate_npg.py index 8015f67b5..6bad6797f 100644 --- a/flox/aggregate_npg.py +++ b/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, diff --git a/flox/aggregations.py b/flox/aggregations.py index 0cf6b259b..d7db7e0fd 100644 --- a/flox/aggregations.py +++ b/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( diff --git a/setup.cfg b/setup.cfg index 3645e5bc7..9dfeb4cc6 100644 --- a/setup.cfg +++ b/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]