ENH: Adding keepdims to np.argmin,np.argmax

doc/release/upcoming_changes/19211.new_feature.rst

@@ -0,0 +1,9 @@
+`keepdims` optional argument added to `numpy.argmin`, `numpy.argmax`
+--------------------------------------------------------------------
+
+`keepdims` argument has been added to `numpy.argmin`, `numpy.argmax`.
+By default, it is `False` and hence no behaviour change should be expected
+in existing codes using `numpy.argmin`, `numpy.argmax`. If set to `True`,
+the axes which are reduced are left in the result as dimensions with size one.
+In simple words, the resulting array will be having same dimensions
+as the input array.

doc/source/reference/c-api/array.rst

@@ -2014,7 +2014,8 @@ Calculation
 
 
 .. c:function:: PyObject* PyArray_ArgMax( \
-        PyArrayObject* self, int axis, PyArrayObject* out)
+        PyArrayObject* self, int axis, PyArrayObject* out, \
+        int keepdims)
 
     Equivalent to :meth:`ndarray.argmax<numpy.ndarray.argmax>` (*self*, *axis*). Return the index of
     the largest element of *self* along *axis*.

numpy/__init__.cython-30.pxd

@@ -643,7 +643,7 @@ cdef extern from "numpy/arrayobject.h":
     int PyArray_Sort (ndarray, int, NPY_SORTKIND)
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
-    object PyArray_ArgMax (ndarray, int, ndarray)
+    object PyArray_ArgMax (ndarray, int, ndarray, int)
     object PyArray_ArgMin (ndarray, int, ndarray)
     object PyArray_Reshape (ndarray, object)
     object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER)

numpy/__init__.pxd

@@ -601,7 +601,7 @@ cdef extern from "numpy/arrayobject.h":
     int PyArray_Sort (ndarray, int, NPY_SORTKIND)
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
-    object PyArray_ArgMax (ndarray, int, ndarray)
+    object PyArray_ArgMax (ndarray, int, ndarray, int)
     object PyArray_ArgMin (ndarray, int, ndarray)
     object PyArray_Reshape (ndarray, object)
     object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER)

numpy/__init__.pyi

@@ -1299,37 +1299,43 @@ class _ArrayOrScalarCommon:
         self,
         axis: None = ...,
         out: None = ...,
+        keepdims: L[False] = ...,
     ) -> intp: ...
     @overload
     def argmax(
         self,
         axis: _ShapeLike = ...,
         out: None = ...,
+        keepdims: bool = ...,
     ) -> Any: ...
     @overload
     def argmax(
         self,
         axis: Optional[_ShapeLike] = ...,
         out: _NdArraySubClass = ...,
+        keepdims: bool = ...,
     ) -> _NdArraySubClass: ...
 
     @overload
     def argmin(
         self,
         axis: None = ...,
         out: None = ...,
+        keepdims: L[False] = ...,
     ) -> intp: ...
     @overload
     def argmin(
         self,
         axis: _ShapeLike = ...,
-         out: None = ...,
+        out: None = ...,
+        keepdims: bool = ...,
     ) -> Any: ...
     @overload
     def argmin(
         self,
         axis: Optional[_ShapeLike] = ...,
         out: _NdArraySubClass = ...,
+        keepdims: bool = ...,
     ) -> _NdArraySubClass: ...
 
     def argsort(

numpy/core/fromnumeric.py

@@ -1114,12 +1114,12 @@ def argsort(a, axis=-1, kind=None, order=None):
     return _wrapfunc(a, 'argsort', axis=axis, kind=kind, order=order)
 
 
-def _argmax_dispatcher(a, axis=None, out=None):
-    return (a, out)
+def _argmax_dispatcher(a, axis=None, out=None, keepdims=None):
+    return (a, out, keepdims)
 
 
 @array_function_dispatch(_argmax_dispatcher)
-def argmax(a, axis=None, out=None):
+def argmax(a, axis=None, out=None, keepdims=False):
     """
     Returns the indices of the maximum values along an axis.
 
@@ -1133,12 +1133,18 @@ def argmax(a, axis=None, out=None):
     out : array, optional
         If provided, the result will be inserted into this array. It should
         be of the appropriate shape and dtype.
+    keepdims : bool, optional
+            If this is set to True, the axes which are reduced are left
+            in the result as dimensions with size one. With this option,
+            the result will broadcast correctly against the array.
 
     Returns
     -------
     index_array : ndarray of ints
         Array of indices into the array. It has the same shape as `a.shape`
-        with the dimension along `axis` removed.
+        with the dimension along `axis` removed. If `keepdims` is set to True,
+        then the size of `axis` will be 1 with the resulting array having same
+        shape as `a.shape`.
 
     See Also
     --------
@@ -1191,16 +1197,23 @@ def argmax(a, axis=None, out=None):
     >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1)
     array([4, 3])
 
+    Setting `keepdims` to `True`,
+
+    >>> x = np.arange(24).reshape((2, 3, 4))
+    >>> res = np.argmax(x, axis=1, keepdims=True)
+    >>> res.shape
+    (2, 1, 4)
     """
-    return _wrapfunc(a, 'argmax', axis=axis, out=out)
+    return _wrapfunc(a, 'argmax', axis=axis, out=out, 
+                        keepdims=keepdims)
 
 
-def _argmin_dispatcher(a, axis=None, out=None):
-    return (a, out)
+def _argmin_dispatcher(a, axis=None, out=None, keepdims=None):
+    return (a, out, keepdims)
 
 
 @array_function_dispatch(_argmin_dispatcher)
-def argmin(a, axis=None, out=None):
+def argmin(a, axis=None, out=None, keepdims=False):
     """
     Returns the indices of the minimum values along an axis.
 
@@ -1214,12 +1227,18 @@ def argmin(a, axis=None, out=None):
     out : array, optional
         If provided, the result will be inserted into this array. It should
         be of the appropriate shape and dtype.
+    keepdims : bool, optional
+        If this is set to True, the axes which are reduced are left
+        in the result as dimensions with size one. With this option,
+        the result will broadcast correctly against the array.
 
     Returns
     -------
     index_array : ndarray of ints
         Array of indices into the array. It has the same shape as `a.shape`
-        with the dimension along `axis` removed.
+        with the dimension along `axis` removed. If `keepdims` is set to True,
+        then the size of `axis` will be 1 with the resulting array having same
+        shape as `a.shape`.
 
     See Also
     --------
@@ -1272,8 +1291,24 @@ def argmin(a, axis=None, out=None):
     >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1)
     array([2, 0])
 
+    Setting `keepdims` to `True`,
+
+    >>> x = np.arange(24).reshape((2, 3, 4))
+    >>> res = np.argmin(x, axis=1, keepdims=True)
+    >>> res.shape
+    (2, 1, 4)
     """
-    return _wrapfunc(a, 'argmin', axis=axis, out=out)
+    res = _wrapfunc(a, 'argmin', axis=axis, out=out)
+
+    if keepdims:
+        if axis is None:
+            new_shape = (1,)*a.ndim
+        else:
+            new_shape = list(a.shape)
+            new_shape[axis] = 1
+        return res.reshape(new_shape)
+
+    return res
 
 
 def _searchsorted_dispatcher(a, v, side=None, sorter=None):

numpy/core/fromnumeric.pyi

@@ -130,25 +130,29 @@ def argmax(
     a: ArrayLike,
     axis: None = ...,
     out: Optional[ndarray] = ...,
+    keepdims: Literal[False] = ...,
 ) -> intp: ...
 @overload
 def argmax(
     a: ArrayLike,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
+    keepdims: bool = ...,
 ) -> Any: ...
 
 @overload
 def argmin(
     a: ArrayLike,
     axis: None = ...,
     out: Optional[ndarray] = ...,
+    keepdims: Literal[False] = ...,
 ) -> intp: ...
 @overload
 def argmin(
     a: ArrayLike,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
+    keepdims: bool = ...,
 ) -> Any: ...
 
 @overload

numpy/core/src/multiarray/calculation.c

@@ -38,7 +38,8 @@ power_of_ten(int n)
  * ArgMax
  */
 NPY_NO_EXPORT PyObject *
-PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out)
+PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out,
+               int keepdims)
 {
     PyArrayObject *ap = NULL, *rp = NULL;
     PyArray_ArgFunc* arg_func;

numpy/core/src/multiarray/calculation.h

@@ -2,7 +2,8 @@
 #define _NPY_CALCULATION_H_
 
 NPY_NO_EXPORT PyObject*
-PyArray_ArgMax(PyArrayObject* self, int axis, PyArrayObject *out);
+PyArray_ArgMax(PyArrayObject* self, int axis, PyArrayObject *out, 
+               int keepdims);
 
 NPY_NO_EXPORT PyObject*
 PyArray_ArgMin(PyArrayObject* self, int axis, PyArrayObject *out);

numpy/core/src/multiarray/methods.c

@@ -284,16 +284,18 @@ array_argmax(PyArrayObject *self,
 {
     int axis = NPY_MAXDIMS;
     PyArrayObject *out = NULL;
+    int keepdims = -1;
     NPY_PREPARE_ARGPARSER;
 
     if (npy_parse_arguments("argmax", args, len_args, kwnames,
             "|axis", &PyArray_AxisConverter, &axis,
             "|out", &PyArray_OutputConverter, &out,
+            "|keepdims", &PyArray_BoolConverter, &keepdims,
             NULL, NULL, NULL) < 0) {
         return NULL;
     }
 
-    PyObject *ret = PyArray_ArgMax(self, axis, out);
+    PyObject *ret = PyArray_ArgMax(self, axis, out, keepdims);
 
     /* this matches the unpacking behavior of ufuncs */
     if (out == NULL) {

numpy/core/tests/test_multiarray.py

@@ -4333,7 +4333,22 @@ def test_object_argmax_with_NULLs(self):
         assert_equal(a.argmax(), 3)
         a[1] = 30
         assert_equal(a.argmax(), 1)
+
+    def test_np_argmax_keepdims(self):
 
+        sizes = [(3,), (3, 2), (2, 3),
+                 (3, 3), (2, 3, 4), (4, 3, 2)]
+        for size in sizes:
+            arr = np.random.normal(size=size)
+            for axis in range(len(size)):
+                res = np.argmax(arr, axis=axis, keepdims=True)
+                assert_(res.ndim == arr.ndim)
+                assert_(res.shape[axis] == 1)
+
+            # Testing for axis=None, keepdims=True
+            res = np.argmin(arr, axis=None, keepdims=True)
+            assert_(res.ndim == arr.ndim)
+            assert_(res.shape == (1,)*arr.ndim)
 
 class TestArgmin:
 
@@ -4489,7 +4504,22 @@ def test_object_argmin_with_NULLs(self):
         assert_equal(a.argmin(), 3)
         a[1] = 10
         assert_equal(a.argmin(), 1)
+
+    def test_np_argmin_keepdims(self):
 
+        sizes = [(3,), (3, 2), (2, 3),
+                 (3, 3), (2, 3, 4), (4, 3, 2)]
+        for size in sizes:
+            arr = np.random.normal(size=size)
+            for axis in range(len(size)):
+                res = np.argmin(arr, axis=axis, keepdims=True)
+                assert_(res.ndim == arr.ndim)
+                assert_(res.shape[axis] == 1)
+
+            # Testing for axis=None, keepdims=True
+            res = np.argmin(arr, axis=None, keepdims=True)
+            assert_(res.ndim == arr.ndim)
+            assert_(res.shape == (1,)*arr.ndim)
 
 class TestMinMax:
 

numpy/ma/core.py

@@ -5491,7 +5491,7 @@ def argsort(self, axis=np._NoValue, kind=None, order=None,
         filled = self.filled(fill_value)
         return filled.argsort(axis=axis, kind=kind, order=order)
 
-    def argmin(self, axis=None, fill_value=None, out=None):
+    def argmin(self, axis=None, fill_value=None, out=None, keepdims=False):
         """
         Return array of indices to the minimum values along the given axis.
 
@@ -5534,9 +5534,9 @@ def argmin(self, axis=None, fill_value=None, out=None):
         if fill_value is None:
             fill_value = minimum_fill_value(self)
         d = self.filled(fill_value).view(ndarray)
-        return d.argmin(axis, out=out)
+        return d.argmin(axis, out=out, keepdims=keepdims)
 
-    def argmax(self, axis=None, fill_value=None, out=None):
+    def argmax(self, axis=None, fill_value=None, out=None, keepdims=False):
         """
         Returns array of indices of the maximum values along the given axis.
         Masked values are treated as if they had the value fill_value.
@@ -5571,7 +5571,7 @@ def argmax(self, axis=None, fill_value=None, out=None):
         if fill_value is None:
             fill_value = maximum_fill_value(self._data)
         d = self.filled(fill_value).view(ndarray)
-        return d.argmax(axis, out=out)
+        return d.argmax(axis, out=out, keepdims=keepdims)
 
     def sort(self, axis=-1, kind=None, order=None,
              endwith=True, fill_value=None):

numpy/ma/core.pyi

@@ -270,8 +270,8 @@ class MaskedArray(ndarray[_ShapeType, _DType_co]):
     def std(self, axis=..., dtype=..., out=..., ddof=..., keepdims=...): ...
     def round(self, decimals=..., out=...): ...
     def argsort(self, axis=..., kind=..., order=..., endwith=..., fill_value=...): ...
-    def argmin(self, axis=..., fill_value=..., out=...): ...
-    def argmax(self, axis=..., fill_value=..., out=...): ...
+    def argmin(self, axis=..., fill_value=..., out=..., keepdims=...): ...
+    def argmax(self, axis=..., fill_value=..., out=..., keepdims=...): ...
     def sort(self, axis=..., kind=..., order=..., endwith=..., fill_value=...): ...
     def min(self, axis=..., out=..., fill_value=..., keepdims=...): ...
     # NOTE: deprecated