diff --git a/doc/release/upcoming_changes/19211.new_feature.rst b/doc/release/upcoming_changes/19211.new_feature.rst new file mode 100644 index 000000000000..40e42387c153 --- /dev/null +++ b/doc/release/upcoming_changes/19211.new_feature.rst @@ -0,0 +1,7 @@ +``keepdims`` optional argument added to `numpy.argmin`, `numpy.argmax` +---------------------------------------------------------------------- + +``keepdims`` argument is added to `numpy.argmin`, `numpy.argmax`. +If set to ``True``, the axes which are reduced are left in the result as dimensions with size one. +The resulting array has the same number of dimensions and will broadcast with the +input array. diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index ec38d49439a0..22929da539ee 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -1299,18 +1299,24 @@ 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 @@ -1318,18 +1324,24 @@ class _ArrayOrScalarCommon: 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( diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index 65a42eb1ee72..9c6af47673c2 100644 --- a/numpy/core/fromnumeric.py +++ b/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): +def _argmax_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): return (a, out) @array_function_dispatch(_argmax_dispatcher) -def argmax(a, axis=None, out=None): +def argmax(a, axis=None, out=None, *, keepdims=np._NoValue): """ 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) + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmax', axis=axis, out=out, **kwds) -def _argmin_dispatcher(a, axis=None, out=None): +def _argmin_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): return (a, out) @array_function_dispatch(_argmin_dispatcher) -def argmin(a, axis=None, out=None): +def argmin(a, axis=None, out=None, *, keepdims=np._NoValue): """ 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,15 @@ 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) + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmin', axis=axis, out=out, **kwds) def _searchsorted_dispatcher(a, v, side=None, sorter=None): diff --git a/numpy/core/fromnumeric.pyi b/numpy/core/fromnumeric.pyi index 3342ec3ac47b..45057e4b1299 100644 --- a/numpy/core/fromnumeric.pyi +++ b/numpy/core/fromnumeric.pyi @@ -130,12 +130,16 @@ 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 @@ -143,12 +147,16 @@ 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 diff --git a/numpy/core/src/multiarray/calculation.c b/numpy/core/src/multiarray/calculation.c index de67b35b53d6..e89018889d10 100644 --- a/numpy/core/src/multiarray/calculation.c +++ b/numpy/core/src/multiarray/calculation.c @@ -34,11 +34,9 @@ power_of_ten(int n) return ret; } -/*NUMPY_API - * ArgMax - */ NPY_NO_EXPORT PyObject * -PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) +_PyArray_ArgMaxWithKeepdims(PyArrayObject *op, + int axis, PyArrayObject *out, int keepdims) { PyArrayObject *ap = NULL, *rp = NULL; PyArray_ArgFunc* arg_func; @@ -46,6 +44,14 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) npy_intp *rptr; npy_intp i, n, m; int elsize; + // Keep a copy because axis changes via call to PyArray_CheckAxis + int axis_copy = axis; + npy_intp _shape_buf[NPY_MAXDIMS]; + npy_intp *out_shape; + // Keep the number of dimensions and shape of + // original array. Helps when `keepdims` is True. + npy_intp* original_op_shape = PyArray_DIMS(op); + int out_ndim = PyArray_NDIM(op); NPY_BEGIN_THREADS_DEF; if ((ap = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0)) == NULL) { @@ -86,6 +92,29 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) if (ap == NULL) { return NULL; } + + // Decides the shape of the output array. + if (!keepdims) { + out_ndim = PyArray_NDIM(ap) - 1; + out_shape = PyArray_DIMS(ap); + } + else { + out_shape = _shape_buf; + if (axis_copy == NPY_MAXDIMS) { + for (int i = 0; i < out_ndim; i++) { + out_shape[i] = 1; + } + } + else { + /* + * While `ap` may be transposed, we can ignore this for `out` because the + * transpose only reorders the size 1 `axis` (not changing memory layout). + */ + memcpy(out_shape, original_op_shape, out_ndim * sizeof(npy_intp)); + out_shape[axis] = 1; + } + } + arg_func = PyArray_DESCR(ap)->f->argmax; if (arg_func == NULL) { PyErr_SetString(PyExc_TypeError, @@ -103,16 +132,16 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) if (!out) { rp = (PyArrayObject *)PyArray_NewFromDescr( Py_TYPE(ap), PyArray_DescrFromType(NPY_INTP), - PyArray_NDIM(ap) - 1, PyArray_DIMS(ap), NULL, NULL, + out_ndim, out_shape, NULL, NULL, 0, (PyObject *)ap); if (rp == NULL) { goto fail; } } else { - if ((PyArray_NDIM(out) != PyArray_NDIM(ap) - 1) || - !PyArray_CompareLists(PyArray_DIMS(out), PyArray_DIMS(ap), - PyArray_NDIM(out))) { + if ((PyArray_NDIM(out) != out_ndim) || + !PyArray_CompareLists(PyArray_DIMS(out), out_shape, + out_ndim)) { PyErr_SetString(PyExc_ValueError, "output array does not match result of np.argmax."); goto fail; @@ -135,7 +164,7 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) NPY_END_THREADS_DESCR(PyArray_DESCR(ap)); Py_DECREF(ap); - /* Trigger the UPDATEIFCOPY/WRTIEBACKIFCOPY if necessary */ + /* Trigger the UPDATEIFCOPY/WRITEBACKIFCOPY if necessary */ if (out != NULL && out != rp) { PyArray_ResolveWritebackIfCopy(rp); Py_DECREF(rp); @@ -151,10 +180,17 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) } /*NUMPY_API - * ArgMin + * ArgMax */ NPY_NO_EXPORT PyObject * -PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) +PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) +{ + return _PyArray_ArgMaxWithKeepdims(op, axis, out, 0); +} + +NPY_NO_EXPORT PyObject * +_PyArray_ArgMinWithKeepdims(PyArrayObject *op, + int axis, PyArrayObject *out, int keepdims) { PyArrayObject *ap = NULL, *rp = NULL; PyArray_ArgFunc* arg_func; @@ -162,6 +198,14 @@ PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) npy_intp *rptr; npy_intp i, n, m; int elsize; + // Keep a copy because axis changes via call to PyArray_CheckAxis + int axis_copy = axis; + npy_intp _shape_buf[NPY_MAXDIMS]; + npy_intp *out_shape; + // Keep the number of dimensions and shape of + // original array. Helps when `keepdims` is True. + npy_intp* original_op_shape = PyArray_DIMS(op); + int out_ndim = PyArray_NDIM(op); NPY_BEGIN_THREADS_DEF; if ((ap = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0)) == NULL) { @@ -202,6 +246,27 @@ PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) if (ap == NULL) { return NULL; } + + // Decides the shape of the output array. + if (!keepdims) { + out_ndim = PyArray_NDIM(ap) - 1; + out_shape = PyArray_DIMS(ap); + } else { + out_shape = _shape_buf; + if (axis_copy == NPY_MAXDIMS) { + for (int i = 0; i < out_ndim; i++) { + out_shape[i] = 1; + } + } else { + /* + * While `ap` may be transposed, we can ignore this for `out` because the + * transpose only reorders the size 1 `axis` (not changing memory layout). + */ + memcpy(out_shape, original_op_shape, out_ndim * sizeof(npy_intp)); + out_shape[axis] = 1; + } + } + arg_func = PyArray_DESCR(ap)->f->argmin; if (arg_func == NULL) { PyErr_SetString(PyExc_TypeError, @@ -219,16 +284,15 @@ PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) if (!out) { rp = (PyArrayObject *)PyArray_NewFromDescr( Py_TYPE(ap), PyArray_DescrFromType(NPY_INTP), - PyArray_NDIM(ap) - 1, PyArray_DIMS(ap), NULL, NULL, + out_ndim, out_shape, NULL, NULL, 0, (PyObject *)ap); if (rp == NULL) { goto fail; } } else { - if ((PyArray_NDIM(out) != PyArray_NDIM(ap) - 1) || - !PyArray_CompareLists(PyArray_DIMS(out), PyArray_DIMS(ap), - PyArray_NDIM(out))) { + if ((PyArray_NDIM(out) != out_ndim) || + !PyArray_CompareLists(PyArray_DIMS(out), out_shape, out_ndim)) { PyErr_SetString(PyExc_ValueError, "output array does not match result of np.argmin."); goto fail; @@ -266,6 +330,15 @@ PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) return NULL; } +/*NUMPY_API + * ArgMin + */ +NPY_NO_EXPORT PyObject * +PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) +{ + return _PyArray_ArgMinWithKeepdims(op, axis, out, 0); +} + /*NUMPY_API * Max */ diff --git a/numpy/core/src/multiarray/calculation.h b/numpy/core/src/multiarray/calculation.h index 34bc31f69806..49105a1385c0 100644 --- a/numpy/core/src/multiarray/calculation.h +++ b/numpy/core/src/multiarray/calculation.h @@ -4,9 +4,15 @@ NPY_NO_EXPORT PyObject* PyArray_ArgMax(PyArrayObject* self, int axis, PyArrayObject *out); +NPY_NO_EXPORT PyObject* +_PyArray_ArgMaxWithKeepdims(PyArrayObject* self, int axis, PyArrayObject *out, int keepdims); + NPY_NO_EXPORT PyObject* PyArray_ArgMin(PyArrayObject* self, int axis, PyArrayObject *out); +NPY_NO_EXPORT PyObject* +_PyArray_ArgMinWithKeepdims(PyArrayObject* self, int axis, PyArrayObject *out, int keepdims); + NPY_NO_EXPORT PyObject* PyArray_Max(PyArrayObject* self, int axis, PyArrayObject* out); diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 251e527a6b96..dc23b3471cb2 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -284,16 +284,18 @@ array_argmax(PyArrayObject *self, { int axis = NPY_MAXDIMS; PyArrayObject *out = NULL; + npy_bool keepdims = NPY_FALSE; 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_ArgMaxWithKeepdims(self, axis, out, keepdims); /* this matches the unpacking behavior of ufuncs */ if (out == NULL) { @@ -310,16 +312,17 @@ array_argmin(PyArrayObject *self, { int axis = NPY_MAXDIMS; PyArrayObject *out = NULL; + npy_bool keepdims = NPY_FALSE; NPY_PREPARE_ARGPARSER; - if (npy_parse_arguments("argmin", 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_ArgMin(self, axis, out); + PyObject *ret = _PyArray_ArgMinWithKeepdims(self, axis, out, keepdims); /* this matches the unpacking behavior of ufuncs */ if (out == NULL) { diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index 25dd76256663..8438c8bd5ddc 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -4192,6 +4192,88 @@ def test_unicode(self): assert_array_equal(g1 < g2, [g1[i] < g2[i] for i in [0, 1, 2]]) assert_array_equal(g1 > g2, [g1[i] > g2[i] for i in [0, 1, 2]]) +class TestArgmaxArgminCommon: + + sizes = [(), (3,), (3, 2), (2, 3), + (3, 3), (2, 3, 4), (4, 3, 2), + (1, 2, 3, 4), (2, 3, 4, 1), + (3, 4, 1, 2), (4, 1, 2, 3)] + + @pytest.mark.parametrize("size, axis", itertools.chain(*[[(size, axis) + for axis in list(range(-len(size), len(size))) + [None]] + for size in sizes])) + @pytest.mark.parametrize('method', [np.argmax, np.argmin]) + def test_np_argmin_argmax_keepdims(self, size, axis, method): + + arr = np.random.normal(size=size) + + # contiguous arrays + if axis is None: + new_shape = [1 for _ in range(len(size))] + else: + new_shape = list(size) + new_shape[axis] = 1 + new_shape = tuple(new_shape) + + _res_orig = method(arr, axis=axis) + res_orig = _res_orig.reshape(new_shape) + res = method(arr, axis=axis, keepdims=True) + assert_equal(res, res_orig) + assert_(res.shape == new_shape) + outarray = np.empty(res.shape, dtype=res.dtype) + res1 = method(arr, axis=axis, out=outarray, + keepdims=True) + assert_(res1 is outarray) + assert_equal(res, outarray) + + if len(size) > 0: + wrong_shape = list(new_shape) + if axis is not None: + wrong_shape[axis] = 2 + else: + wrong_shape[0] = 2 + wrong_outarray = np.empty(wrong_shape, dtype=res.dtype) + with pytest.raises(ValueError): + method(arr.T, axis=axis, + out=wrong_outarray, keepdims=True) + + # non-contiguous arrays + if axis is None: + new_shape = [1 for _ in range(len(size))] + else: + new_shape = list(size)[::-1] + new_shape[axis] = 1 + new_shape = tuple(new_shape) + + _res_orig = method(arr.T, axis=axis) + res_orig = _res_orig.reshape(new_shape) + res = method(arr.T, axis=axis, keepdims=True) + assert_equal(res, res_orig) + assert_(res.shape == new_shape) + outarray = np.empty(new_shape[::-1], dtype=res.dtype) + outarray = outarray.T + res1 = method(arr.T, axis=axis, out=outarray, + keepdims=True) + assert_(res1 is outarray) + assert_equal(res, outarray) + + if len(size) > 0: + # one dimension lesser for non-zero sized + # array should raise an error + with pytest.raises(ValueError): + method(arr[0], axis=axis, + out=outarray, keepdims=True) + + if len(size) > 0: + wrong_shape = list(new_shape) + if axis is not None: + wrong_shape[axis] = 2 + else: + wrong_shape[0] = 2 + wrong_outarray = np.empty(wrong_shape, dtype=res.dtype) + with pytest.raises(ValueError): + method(arr.T, axis=axis, + out=wrong_outarray, keepdims=True) class TestArgmax: diff --git a/numpy/ma/core.py b/numpy/ma/core.py index 82e5e7155322..a3d85d0a199e 100644 --- a/numpy/ma/core.py +++ b/numpy/ma/core.py @@ -5491,7 +5491,8 @@ 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=np._NoValue): """ Return array of indices to the minimum values along the given axis. @@ -5534,9 +5535,11 @@ 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) + keepdims = False if keepdims is np._NoValue else bool(keepdims) + 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=np._NoValue): """ 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 +5574,8 @@ 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) + keepdims = False if keepdims is np._NoValue else bool(keepdims) + return d.argmax(axis, out=out, keepdims=keepdims) def sort(self, axis=-1, kind=None, order=None, endwith=True, fill_value=None): diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi index e7e3f1f36818..bc1f45a8d5ad 100644 --- a/numpy/ma/core.pyi +++ b/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