Fixes for latest numpy_groupies (#85)

* Remove nanvar, nanstd compatibility code

* Remove argreduction compatibility code

* Fix argreduction along axis with single chunk

Fix argreductions

* Update min numpy_groupies version

* Restore tests

* Fix numpy tests.

* Skip nanagr* instead of xfail

Author: dcherian

ci/environment.yml

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

ci/minimal-requirements.yml

@@ -8,7 +8,7 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-xdist
-  - numpy_groupies
+  - numpy_groupies>=0.9.15
   - pandas
   - pooch
   - toolz

ci/no-dask.yml

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

ci/no-xarray.yml

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

flox/aggregations.py

@@ -9,6 +9,14 @@
 from . import aggregate_flox, aggregate_npg, xrdtypes as dtypes, xrutils
 
 
+def _is_arg_reduction(func: str | Aggregation) -> bool:
+    if isinstance(func, str) and func in ["argmin", "argmax", "nanargmax", "nanargmin"]:
+        return True
+    if isinstance(func, Aggregation) and func.reduction_type == "argreduce":
+        return True
+    return False
+
+
 def generic_aggregate(
     group_idx,
     array,
@@ -488,7 +496,11 @@ def _initialize_aggregation(
     agg.fill_value[func] = _get_fill_value(agg.dtype[func], agg.fill_value[func])
 
     fv = fill_value if fill_value is not None else agg.fill_value[agg.name]
-    agg.fill_value["numpy"] = (fv,)
+    if _is_arg_reduction(agg):
+        # this allows us to unravel_index easily. we have to do that nearly every time.
+        agg.fill_value["numpy"] = (0,)
+    else:
+        agg.fill_value["numpy"] = (fv,)
 
     if finalize_kwargs is not None:
         assert isinstance(finalize_kwargs, dict)

flox/core.py

@@ -509,6 +509,7 @@ def chunk_argreduce(
     dask.array.reductions.argtopk
     """
     array, idx = array_plus_idx
+    by = np.broadcast_to(by, array.shape)
 
     results = chunk_reduce(
         array,
@@ -522,17 +523,22 @@ def chunk_argreduce(
         sort=sort,
     )
     if not isnull(results["groups"]).all():
-        # will not work for empty groups...
-        # glorious
         idx = np.broadcast_to(idx, array.shape)
+
+        # array, by get flattened to 1D before passing to npg
+        # so the indexes need to be unraveled
         newidx = np.unravel_index(results["intermediates"][1], array.shape)
+
+        # Now index into the actual "global" indexes `idx`
         results["intermediates"][1] = idx[newidx]
 
     if reindex and expected_groups is not None:
         results["intermediates"][1] = reindex_(
             results["intermediates"][1], results["groups"].squeeze(), expected_groups, fill_value=0
         )
 
+    assert results["intermediates"][0].shape == results["intermediates"][1].shape
+
     return results
 
 
@@ -879,34 +885,45 @@ def _grouped_combine(
         array_idx = tuple(
             _conc2(x_chunk, key1="intermediates", key2=idx, axis=axis) for idx in (0, 1)
         )
-        results = chunk_argreduce(
-            array_idx,
-            groups,
-            func=agg.combine[slicer],  # count gets treated specially next
-            axis=axis,
-            expected_groups=None,
-            fill_value=agg.fill_value["intermediate"][slicer],
-            dtype=agg.dtype["intermediate"][slicer],
-            engine=engine,
-            sort=sort,
-        )
+
+        # for a single element along axis, we don't want to run the argreduction twice
+        # This happens when we are reducing along an axis with a single chunk.
+        avoid_reduction = array_idx[0].shape[axis[0]] == 1
+        if avoid_reduction:
+            results = {"groups": groups, "intermediates": list(array_idx)}
+        else:
+            results = chunk_argreduce(
+                array_idx,
+                groups,
+                func=agg.combine[slicer],  # count gets treated specially next
+                axis=axis,
+                expected_groups=None,
+                fill_value=agg.fill_value["intermediate"][slicer],
+                dtype=agg.dtype["intermediate"][slicer],
+                engine=engine,
+                sort=sort,
+            )
 
         if agg.chunk[-1] == "nanlen":
             counts = _conc2(x_chunk, key1="intermediates", key2=2, axis=axis)
-            # sum the counts
-            results["intermediates"].append(
-                chunk_reduce(
-                    counts,
-                    groups,
-                    func="sum",
-                    axis=axis,
-                    expected_groups=None,
-                    fill_value=(0,),
-                    dtype=(np.intp,),
-                    engine=engine,
-                    sort=sort,
-                )["intermediates"][0]
-            )
+
+            if avoid_reduction:
+                results["intermediates"].append(counts)
+            else:
+                # sum the counts
+                results["intermediates"].append(
+                    chunk_reduce(
+                        counts,
+                        groups,
+                        func="sum",
+                        axis=axis,
+                        expected_groups=None,
+                        fill_value=(0,),
+                        dtype=(np.intp,),
+                        engine=engine,
+                        sort=sort,
+                    )["intermediates"][0]
+                )
 
     elif agg.reduction_type == "reduce":
         # Here we reduce the intermediates individually
@@ -1006,24 +1023,7 @@ def _reduce_blockwise(array, by, agg, *, axis, expected_groups, fill_value, engi
     )  # type: ignore
 
     if _is_arg_reduction(agg):
-        if array.ndim > 1:
-            # default fill_value is -1; we can't unravel that;
-            # so replace -1 with 0; unravel; then replace 0 with -1
-            # UGH!
-            idx = results["intermediates"][0]
-            mask = idx == agg.fill_value["numpy"][0]
-            idx[mask] = 0
-            # Fix npg bug where argmax with nD array, 1D group_idx, axis=-1
-            # will return wrong indices
-            idx = np.unravel_index(idx, array.shape)[-1]
-            idx[mask] = agg.fill_value["numpy"][0]
-            results["intermediates"][0] = idx
-    elif agg.name in ["nanvar", "nanstd"]:
-        # TODO: Fix npg bug where all-NaN rows are 0 instead of NaN
-        value, counts = results["intermediates"]
-        mask = counts <= 0
-        value[mask] = np.nan
-        results["intermediates"][0] = value
+        results["intermediates"][0] = np.unravel_index(results["intermediates"][0], array.shape)[-1]
 
     result = _finalize_results(
         results, agg, axis, expected_groups, fill_value=fill_value, reindex=reindex
@@ -1530,12 +1530,7 @@ def groupby_reduce(
     #     The only way to do this consistently is mask out using min_count
     #     Consider np.sum([np.nan]) = np.nan, np.nansum([np.nan]) = 0
     if min_count is None:
-        if (
-            len(axis) < by.ndim
-            or fill_value is not None
-            # TODO: Fix npg bug where all-NaN rows are 0 instead of NaN
-            or (not has_dask and isinstance(func, str) and func in ["nanvar", "nanstd"])
-        ):
+        if len(axis) < by.ndim or fill_value is not None:
             min_count = 1
 
     # TODO: set in xarray?

setup.cfg

@@ -28,7 +28,7 @@ include_package_data = True
 python_requires = >=3.7
 install_requires =
     pandas
-    numpy_groupies
+    numpy_groupies >= '0.9.15'
     toolz
     importlib-metadata; python_version < '3.8'
 

tests/test_core.py

@@ -53,9 +53,9 @@ def dask_array_ones(*args):
     "min",
     "nanmin",
     "argmax",
-    pytest.param("nanargmax", marks=(pytest.mark.xfail,)),
+    pytest.param("nanargmax", marks=(pytest.mark.skip,)),
     "argmin",
-    pytest.param("nanargmin", marks=(pytest.mark.xfail,)),
+    pytest.param("nanargmin", marks=(pytest.mark.skip,)),
     "any",
     "all",
     pytest.param("median", marks=(pytest.mark.skip,)),
@@ -142,7 +142,7 @@ def gen_array_by(size, func):
     return array, by
 
 
-@pytest.mark.parametrize("chunks", [None, 3, 4])
+@pytest.mark.parametrize("chunks", [None, -1, 3, 4])
 @pytest.mark.parametrize("nby", [1, 2, 3])
 @pytest.mark.parametrize("size", ((12,), (12, 9)))
 @pytest.mark.parametrize("add_nan_by", [True, False])