Repository: numpy/numpy

Issue:
DOC: argmin, argmax, argort, etc. should link to `np.take_along_axis` in their documentation
In #6075, I wondered whether it might be possible to

> to have some simple way to use the output of `np.arg[min|max|sort]` as a proper index (to avoid hackery as in https://github.com/astropy/astropy/blob/master/astropy/time/core.py#L896 -- if there is a simpler way already, let me know!)

and @jaimefrio suggested that it might be simpler to have

> a `make_me_an_indexing_tuple=False` new keyword argument for the `argxxx` functions?

That indeed seems better, so I'm raising this as a separate issue.  The code linked to above could be used as a template on how to generate the index tuple.



Discussion:
Indeed, this problem has also flummoxed me.

Here's the helper function I wrote, which may be slightly more interpretable:

``` python
def argsort_indices(a, axis=-1):
    """Like argsort, but returns an index suitable for sorting the
    the original array even if that array is multidimensional
    """
    a = np.asarray(a)
    ind = list(np.ix_(*[np.arange(d) for d in a.shape]))
    ind[axis] = a.argsort(axis)
    return tuple(ind)
```

@shoyer - yours is definitely easier to read, in that it leaves the proper shaping to `np.ix_`... also, it doesn't deal with `keepdims`, right?

> also, it doesn't deal with keepdims, right?

I guess it's less general in that it only works for `argsort` -- I didn't think about `min`/`max`, so I didn't need anything for `keepdims`.

I think it would be lot better than today's implementation if the `axis=None` option returned something equivalent to :
```python
np.unravel_index(np.argsort(a, axis=None), a.shape)
```
> I think it would be lot better than today's implementation if the axis=None option returned something equivalent to :
> np.unravel_index(np.argsort(a, axis=None), a.shape)

As a rule, NumPy doesn't change existing behavior. So this will need to be new API of some sort.
@shoyer Oh that's true.
I've just submitted a pull request to add my proposition to the doc so that the user can easily find how to retrieve indices.
I think this issue is now irrelevant, since we have `np.take_along_axis` instead, which solves the issue. I will however transform it into a Documentation issue, since those functions should list it under the See Also section at least.

Resolving patch:
diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -796,7 +796,9 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
     --------
     partition : Describes partition algorithms used.
     ndarray.partition : Inplace partition.
-    argsort : Full indirect sort
+    argsort : Full indirect sort.
+    take_along_axis : Apply ``index_array`` from argpartition 
+                      to an array as if by calling partition.
 
     Notes
     -----
@@ -816,6 +818,14 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
     >>> np.array(x)[np.argpartition(x, 3)]
     array([2, 1, 3, 4])
 
+    Multi-dimensional array:
+
+    >>> x = np.array([[3, 4, 2], [1, 3, 1]])
+    >>> index_array = np.argpartition(x, kth=1, axis=-1)
+    >>> np.take_along_axis(x, index_array, axis=-1)  # same as np.partition(x, kth=1)
+    array([[2, 3, 4],
+           [1, 1, 3]])
+
     """
     return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)
 
@@ -1025,6 +1035,8 @@ def argsort(a, axis=-1, kind=None, order=None):
     lexsort : Indirect stable sort with multiple keys.
     ndarray.sort : Inplace sort.
     argpartition : Indirect partial sort.
+    take_along_axis : Apply ``index_array`` from argsort 
+                      to an array as if by calling sort.
 
     Notes
     -----
@@ -1120,6 +1132,8 @@ def argmax(a, axis=None, out=None):
     ndarray.argmax, argmin
     amax : The maximum value along a given axis.
     unravel_index : Convert a flat index into an index tuple.
+    take_along_axis : Apply ``np.expand_dims(index_array, axis)`` 
+                      from argmax to an array as if by calling max.
 
     Notes
     -----
@@ -1154,6 +1168,16 @@ def argmax(a, axis=None, out=None):
     >>> np.argmax(b)  # Only the first occurrence is returned.
     1
 
+    >>> x = np.array([[4,2,3], [1,0,3]])
+    >>> index_array = np.argmax(x, axis=-1)
+    >>> # Same as np.max(x, axis=-1, keepdims=True)
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)
+    array([[4],
+           [3]])
+    >>> # Same as np.max(x, axis=-1)
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1)
+    array([4, 3])
+
     """
     return _wrapfunc(a, 'argmax', axis=axis, out=out)
 
@@ -1189,6 +1213,8 @@ def argmin(a, axis=None, out=None):
     ndarray.argmin, argmax
     amin : The minimum value along a given axis.
     unravel_index : Convert a flat index into an index tuple.
+    take_along_axis : Apply ``np.expand_dims(index_array, axis)`` 
+                      from argmin to an array as if by calling min.
 
     Notes
     -----
@@ -1223,6 +1249,16 @@ def argmin(a, axis=None, out=None):
     >>> np.argmin(b)  # Only the first occurrence is returned.
     0
 
+    >>> x = np.array([[4,2,3], [1,0,3]])
+    >>> index_array = np.argmin(x, axis=-1)
+    >>> # Same as np.min(x, axis=-1, keepdims=True)
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)
+    array([[2],
+           [0]])
+    >>> # Same as np.max(x, axis=-1)
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1)
+    array([2, 0])
+
     """
     return _wrapfunc(a, 'argmin', axis=axis, out=out)
 


Test patch: