Repository: ml31415/numpy-groupies

Issue:
bug with argmax and nd array
``` python
from numpy_groupies.aggregate_numpy import aggregate
import numpy as np

labels = np.array([0, 0, 2, 2, 2, 1, 1, 2, 2, 1, 1, 0], dtype=int)
array = np.array([[1] * 12, [1] * 12])
aggregate(labels, array, axis=-1, func="argmax")
```

gives
```
array([[ 0,  5,  2],
       [12, 17, 14]])
```

The real answer is
```
array([[0, 5, 2],
       [0, 5, 2]])
```
so it looks like a `ravel_multi_index` is missing.

Discussion:

Patch:
diff --git a/numpy_groupies/aggregate_numba.py b/numpy_groupies/aggregate_numba.py
index 7cbc441..3113fa0 100644
--- a/numpy_groupies/aggregate_numba.py
+++ b/numpy_groupies/aggregate_numba.py
@@ -38,8 +38,8 @@ class AggregateOp(object):
 
     def __call__(self, group_idx, a, size=None, fill_value=0, order='C',
                  dtype=None, axis=None, ddof=0):
-        iv = input_validation(group_idx, a, size=size, order=order, axis=axis, check_bounds=False)
-        group_idx, a, flat_size, ndim_idx, size = iv
+        iv = input_validation(group_idx, a, size=size, order=order, axis=axis, check_bounds=False, func=self.func)
+        group_idx, a, flat_size, ndim_idx, size, unravel_shape = iv
 
         # TODO: The typecheck should be done by the class itself, not by check_dtype
         dtype = check_dtype(dtype, self.func, a, len(group_idx))
@@ -57,10 +57,13 @@ class AggregateOp(object):
         self._finalize(ret, counter, fill_value)
 
         if self.outer:
-            return outer
+            ret = outer
 
         # Deal with ndimensional indexing
         if ndim_idx > 1:
+            if unravel_shape is not None:
+                # argreductions only
+                ret = np.unravel_index(ret, unravel_shape)[axis]
             ret = ret.reshape(size, order=order)
         return ret
 
@@ -192,7 +195,7 @@ class AggregateGeneric(AggregateOp):
     def __call__(self, group_idx, a, size=None, fill_value=0, order='C',
                  dtype=None, axis=None, ddof=0):
         iv = input_validation(group_idx, a, size=size, order=order, axis=axis, check_bounds=False)
-        group_idx, a, flat_size, ndim_idx, size = iv
+        group_idx, a, flat_size, ndim_idx, size, _ = iv
 
         # TODO: The typecheck should be done by the class itself, not by check_dtype
         dtype = check_dtype(dtype, self.func, a, len(group_idx))
diff --git a/numpy_groupies/aggregate_numpy.py b/numpy_groupies/aggregate_numpy.py
index e5b6d72..af8f3dd 100644
--- a/numpy_groupies/aggregate_numpy.py
+++ b/numpy_groupies/aggregate_numpy.py
@@ -152,7 +152,7 @@ def _var(group_idx, a, size, fill_value, dtype=np.dtype(np.float64),
         raise ValueError("cannot take variance with scalar a")
     counts = np.bincount(group_idx, minlength=size)
     sums = np.bincount(group_idx, weights=a, minlength=size)
-    with np.errstate(divide='ignore'):
+    with np.errstate(divide='ignore', invalid='ignore'):
         means = sums.astype(dtype) / counts
         ret = np.bincount(group_idx, (a - means[group_idx]) ** 2,
                           minlength=size) / (counts - ddof)
@@ -253,8 +253,9 @@ _impl_dict.update(('nan' + k, v) for k, v in list(_impl_dict.items())
 def _aggregate_base(group_idx, a, func='sum', size=None, fill_value=0,
                     order='C', dtype=None, axis=None, _impl_dict=_impl_dict,
                     _nansqueeze=False, cache=None, **kwargs):
-    group_idx, a, flat_size, ndim_idx, size = input_validation(group_idx, a,
-                                                               size=size, order=order, axis=axis)
+    group_idx, a, flat_size, ndim_idx, size, unravel_shape = input_validation(group_idx, a,
+                                                                              size=size, order=order, axis=axis, func=func)
+
     if group_idx.dtype == np.dtype("uint64"):
         # Force conversion to signed int, to avoid issues with bincount etc later
         group_idx = group_idx.astype(int)
@@ -282,6 +283,8 @@ def _aggregate_base(group_idx, a, func='sum', size=None, fill_value=0,
 
     # deal with ndimensional indexing
     if ndim_idx > 1:
+        if unravel_shape is not None:
+            ret = np.unravel_index(ret, unravel_shape)[axis]
         ret = ret.reshape(size, order=order)
     return ret
 
diff --git a/numpy_groupies/utils_numpy.py b/numpy_groupies/utils_numpy.py
index a15945b..41ec2c9 100644
--- a/numpy_groupies/utils_numpy.py
+++ b/numpy_groupies/utils_numpy.py
@@ -236,7 +236,7 @@ def offset_labels(group_idx, inshape, axis, order, size):
 
 
 def input_validation(group_idx, a, size=None, order='C', axis=None,
-                     ravel_group_idx=True, check_bounds=True, method="ravel"):
+                     ravel_group_idx=True, check_bounds=True, method="ravel", func=None):
     """ Do some fairly extensive checking of group_idx and a, trying to
     give the user as much help as possible with what is wrong. Also,
     convert ndim-indexing to 1d indexing.
@@ -276,10 +276,18 @@ def input_validation(group_idx, a, size=None, order='C', axis=None,
             raise NotImplementedError("when using axis arg, size must be"
                                       "None or scalar.")
         else:
+            is_form_3 = group_idx.ndim == 1 and a.ndim > 1 and axis is not None
+            orig_shape = a.shape if is_form_3 else group_idx.shape
+            if "arg" in func:
+                unravel_shape = orig_shape
+            else:
+                unravel_shape = None
+
             group_idx, size = _ravel_group_idx(group_idx, a, axis, size, order, method=method)
             flat_size = np.prod(size)
             ndim_idx = ndim_a
-            return group_idx.ravel(), a.ravel(), flat_size, ndim_idx, size
+            size = orig_shape if is_form_3 and not callable(func) and "cum" in func else size
+            return group_idx.ravel(), a.ravel(), flat_size, ndim_idx, size, unravel_shape
 
     if ndim_idx == 1:
         if size is None:
@@ -310,7 +318,7 @@ def input_validation(group_idx, a, size=None, order='C', axis=None,
         raise ValueError("group_idx and a must be of the same length, or a"
                          " can be scalar")
 
-    return group_idx, a, flat_size, ndim_idx, size
+    return group_idx, a, flat_size, ndim_idx, size, None
 
 
 ### General tools ###


Test patch:
diff --git a/numpy_groupies/tests/__init__.py b/numpy_groupies/tests/__init__.py
index 3ed1972..9b69489 100644
--- a/numpy_groupies/tests/__init__.py
+++ b/numpy_groupies/tests/__init__.py
@@ -53,3 +53,8 @@ def _wrap_notimplemented_xfail(impl, name=None):
     else:
         _try_xfail.__name__ = impl.__name__
     return _try_xfail
+
+
+func_list = ('sum', 'prod', 'min', 'max', 'all', 'any', 'mean', 'std', 'var', 'len',
+             'argmin', 'argmax', 'anynan', 'allnan', 'cumsum',
+             'nansum', 'nanprod', 'nanmin', 'nanmax', 'nanmean', 'nanstd', 'nanvar','nanlen')
diff --git a/numpy_groupies/tests/test_compare.py b/numpy_groupies/tests/test_compare.py
index e0d9d2d..27514a5 100644
--- a/numpy_groupies/tests/test_compare.py
+++ b/numpy_groupies/tests/test_compare.py
@@ -10,7 +10,7 @@ import pytest
 
 from . import (aggregate_purepy, aggregate_numpy_ufunc, aggregate_numpy,
                aggregate_weave, aggregate_numba, aggregate_pandas,
-               _wrap_notimplemented_xfail, _impl_name)
+               _wrap_notimplemented_xfail, _impl_name, func_list)
 
 class AttrDict(dict):
     __getattr__ = dict.__getitem__
@@ -71,13 +71,6 @@ def func_preserve_order(iterator):
     for i, x in enumerate(iterator, 1):
         tmp += x ** i
     return tmp
-
-
-func_list = ('sum', 'prod', 'min', 'max', 'all', 'any', 'mean', 'std', 'var', 'len',
-             'argmin', 'argmax', 'anynan', 'allnan', 'cumsum', func_arbitrary, func_preserve_order,
-             'nansum', 'nanprod', 'nanmin', 'nanmax', 'nanmean', 'nanstd', 'nanvar','nanlen')
-
-
 @pytest.mark.parametrize(["func", "fill_value"], product(func_list, [0, 1, np.nan]),
                          ids=lambda x: getattr(x, '__name__', x))
 def test_cmp(aggregate_cmp, func, fill_value, decimal=10):
diff --git a/numpy_groupies/tests/test_generic.py b/numpy_groupies/tests/test_generic.py
index cc201d5..3b27db9 100644
--- a/numpy_groupies/tests/test_generic.py
+++ b/numpy_groupies/tests/test_generic.py
@@ -4,7 +4,7 @@ import itertools
 import numpy as np
 import pytest
 
-from . import _implementations, _impl_name, _wrap_notimplemented_xfail
+from . import _implementations, _impl_name, _wrap_notimplemented_xfail, func_list
 
 
 @pytest.fixture(params=_implementations, ids=_impl_name)
@@ -294,3 +294,69 @@ def test_sort(aggregate_all, order):
 
     res = aggregate_all(group_idx, a, func="sort", reverse=reverse)
     np.testing.assert_array_equal(res, ref)
+
+
+@pytest.mark.parametrize("axis", (0, 1))
+@pytest.mark.parametrize("size", ((12,), (12, 5)))
+@pytest.mark.parametrize("func", func_list)
+def test_agg_along_axis(aggregate_all, size, func, axis):
+    if len(size) == axis:
+        pytest.skip()
+
+    group_idx = np.zeros(size[axis], dtype=int)
+    array = np.random.randn(*size)
+
+    # add some NaNs to test out nan-skipping
+    if "nan" in func and "nanarg" not in func:
+        array[[1, 4, 5], ...] = np.nan
+    elif "nanarg" in func and array.ndim > 1:
+        array[[1, 4, 5], 1] = np.nan
+    if func in ["any", "all"]:
+        array = array > 0.5
+
+    # construct expected values for all cases
+    if func == "len":
+        expected = np.array(size[axis])
+    elif func == "nanlen":
+        expected = np.array((~np.isnan(array)).sum(axis=axis))
+    elif func == "anynan":
+        expected = np.isnan(array).any(axis=axis)
+    elif func == "allnan":
+        expected = np.isnan(array).all(axis=axis)
+    else:
+        with np.errstate(invalid="ignore", divide="ignore"):
+            expected = getattr(np, func)(array, axis=axis)
+
+    # The default fill_value is 0, the following makes the output
+    # match numpy
+    fill_values = {
+        "nanprod": 1,
+        "nanvar": np.nan,
+        "nanstd": np.nan,
+        "nanmax": np.nan,
+        "nanmin": np.nan,
+        "nanmean": np.nan,
+    }
+
+    # TODO: These xfails need to be fixed...
+    if len(size) > 1:
+        if func == "allnan":
+            # This only fails for numba
+            pytest.xfail()
+
+    actual = aggregate_all(
+        group_idx, array, axis=axis, func=func, fill_value=fill_values.get(func, 0)
+    )
+    assert actual.ndim == array.ndim
+
+    # argmin, argmax don't support keepdims so we can't use that to construct expected
+    # instead we squeeze out the extra dims in actual.
+    np.testing.assert_allclose(actual.squeeze(), expected)
+
+def test_argreduction_nD_array_1D_idx(aggregate_all):
+    # regression test for GH41
+    labels = np.array([0, 0, 2, 2, 2, 1, 1, 2, 2, 1, 1, 0], dtype=int)
+    array = np.array([[1] * 12, [1] * 12])
+    actual = aggregate_all(labels, array, axis=-1, func="argmax")
+    expected = np.array([[0, 5, 2], [0, 5, 2]])
+    np.testing.assert_equal(actual, expected)