Merge behaviour of indices_min/indices_max into argmin/argmax

When argmin or argmax are called with a sequence for 'dim', they now
return a dict with the indices for each dimension in dim.

Author: johnomotani

doc/api.rst

@@ -366,8 +366,6 @@ Computation
 :py:attr:`~DataArray.argmin`
 :py:attr:`~DataArray.idxmax`
 :py:attr:`~DataArray.idxmin`
-:py:attr:`~DataArray.indices_max`
-:py:attr:`~DataArray.indices_min`
 :py:attr:`~DataArray.max`
 :py:attr:`~DataArray.mean`
 :py:attr:`~DataArray.median`

doc/whats-new.rst

@@ -29,9 +29,11 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
-- Added :py:meth:`DataArray.indices_min` and :py:meth:`DataArray.indices_max`
-  to get a dict of the indices for each dimension of the minimum or maximum of
-  a DataArray. (:pull:`3936`)
+- :py:meth:`DataArray.argmin` and :py:meth:`DataArray.argmax` now support
+  sequences of 'dim' arguments, and if a sequence is passed return a dict
+  (which can be passed to :py:meth:`isel` to get the value of the minimum) of
+  the indices for each dimension of the minimum or maximum of a DataArray.
+  (:pull:`3936`)
   By `John Omotani <https://github.com/johnomotani>`_, thanks to `Keisuke Fujii
   <https://github.com/fujiisoup>`_ for work in :pull:`1469`.
 - Added :py:meth:`DataArray.polyfit` and :py:func:`xarray.polyval` for fitting polynomials. (:issue:`3349`)

xarray/core/dataarray.py

@@ -3730,16 +3730,35 @@ def _unravel_argminmax(
         self,
         argminmax: Hashable,
         dim: Union[Hashable, Sequence[Hashable], None],
-        keep_attrs: bool,
+        axis: Union[int, None],
+        keep_attrs: Optional[bool],
         skipna: Optional[bool],
     ) -> Dict[Hashable, "DataArray"]:
         """Apply argmin or argmax over one or more dimensions, returning the result as a
         dict of DataArray that can be passed directly to isel.
         """
-        if dim is None:
+        if dim is None and axis is None:
+            warnings.warn(
+                "Behaviour of argmin/argmax with neither dim nor axis argument will "
+                "change to return a dict of indices of each dimension. To get a "
+                "single, flat index, please use np.argmin(da) or np.argmax(da) instead "
+                "of da.argmin() or da.argmax().",
+                DeprecationWarning,
+            )
+        if dim is ...:
+            # In future, should do this also when (dim is None and axis is None)
             dim = self.dims
-        if not isinstance(dim, Sequence) or isinstance(dim, str):
-            dim = (dim,)
+        if (
+            dim is None
+            or axis is not None
+            or not isinstance(dim, Sequence)
+            or isinstance(dim, str)
+        ):
+            # Return int index if single dimension is passed, and is not part of a
+            # sequence
+            return getattr(self, str(argminmax))(
+                dim=dim, axis=axis, keep_attrs=keep_attrs, skipna=skipna
+            )
 
         # Get a name for the new dimension that does not conflict with any existing
         # dimension
@@ -3771,9 +3790,10 @@ def _unravel_argminmax(
 
         return result
 
-    def indices_min(
+    def argmin(
         self,
         dim: Union[Hashable, Sequence[Hashable]] = None,
+        axis: Union[int, None] = None,
         keep_attrs: bool = None,
         skipna: bool = None,
     ) -> Dict[Hashable, "DataArray"]:
@@ -3788,6 +3808,9 @@ def indices_min(
         dim : hashable or sequence of hashable, optional
             The dimensions over which to find the minimum. By default, finds minimum over
             all dimensions.
+        axis : int, optional
+            Axis over which to apply `argmin`. Only one of the 'dim' and 'axis' arguments
+            can be supplied.
         keep_attrs : bool, optional
             If True, the attributes (`attrs`) will be copied from the original
             object to the new one.  If False (default), the new object will be
@@ -3815,10 +3838,10 @@ def indices_min(
         >>> array.argmin()
         <xarray.DataArray ()>
         array(2)
-        >>> array.indices_min()
+        >>> array.argmin(...)
         {'x': <xarray.DataArray ()>
         array(2)}
-        >>> array.isel(array.indices_min())
+        >>> array.isel(array.argmin(...))
         array(-1)
 
         >>> array = xr.DataArray([[[3, 2, 1], [3, 1, 2], [2, 1, 3]],
@@ -3836,7 +3859,7 @@ def indices_min(
                [1, 1, 1],
                [0, 0, 1]])
         Dimensions without coordinates: y, z
-        >>> array.indices_min(dim="x")
+        >>> array.argmin(dim=["x"])
         {'x': <xarray.DataArray (y: 3, z: 3)>
         array([[1, 0, 0],
                [1, 1, 1],
@@ -3846,22 +3869,23 @@ def indices_min(
         <xarray.DataArray (y: 3)>
         array([ 1, -5,  1])
         Dimensions without coordinates: y
-        >>> array.indices_min(dim=("x", "z"))
+        >>> array.argmin(dim=["x", "z"])
         {'x': <xarray.DataArray (y: 3)>
         array([0, 1, 0])
         Dimensions without coordinates: y, 'z': <xarray.DataArray (y: 3)>
         array([2, 1, 1])
         Dimensions without coordinates: y}
-        >>> array.isel(array.indices_min(dim=("x", "z")))
+        >>> array.isel(array.argmin(dim=["x", "z"]))
         <xarray.DataArray (y: 3)>
         array([ 1, -5,  1])
         Dimensions without coordinates: y
         """
-        return self._unravel_argminmax("argmin", dim, keep_attrs, skipna)
+        return self._unravel_argminmax("_argmin_base", dim, axis, keep_attrs, skipna)
 
-    def indices_max(
+    def argmax(
         self,
         dim: Union[Hashable, Sequence[Hashable]] = None,
+        axis: Union[int, None] = None,
         keep_attrs: bool = None,
         skipna: bool = None,
     ) -> Dict[Hashable, "DataArray"]:
@@ -3876,6 +3900,9 @@ def indices_max(
         dim : hashable or sequence of hashable, optional
             The dimensions over which to find the maximum. By default, finds maximum over
             all dimensions.
+        axis : int, optional
+            Axis over which to apply `argmin`. Only one of the 'dim' and 'axis' arguments
+            can be supplied.
         keep_attrs : bool, optional
             If True, the attributes (`attrs`) will be copied from the original
             object to the new one.  If False (default), the new object will be
@@ -3903,10 +3930,10 @@ def indices_max(
         >>> array.argmax()
         <xarray.DataArray ()>
         array(3)
-        >>> array.indices_max()
+        >>> array.argmax(...)
         {'x': <xarray.DataArray ()>
         array(3)}
-        >>> array.isel(array.indices_max())
+        >>> array.isel(array.argmax(...))
         <xarray.DataArray ()>
         array(3)
 
@@ -3925,7 +3952,7 @@ def indices_max(
                [0, 1, 0],
                [0, 1, 0]])
         Dimensions without coordinates: y, z
-        >>> array.indices_max(dim="x")
+        >>> array.argmax(dim=["x"])
         {'x': <xarray.DataArray (y: 3, z: 3)>
         array([[0, 1, 1],
                [0, 1, 0],
@@ -3935,18 +3962,18 @@ def indices_max(
         <xarray.DataArray (y: 3)>
         array([3, 5, 3])
         Dimensions without coordinates: y
-        >>> array.indices_max(dim=("x", "z"))
+        >>> array.argmax(dim=["x", "z"])
         {'x': <xarray.DataArray (y: 3)>
         array([0, 1, 0])
         Dimensions without coordinates: y, 'z': <xarray.DataArray (y: 3)>
         array([0, 1, 2])
         Dimensions without coordinates: y}
-        >>> array.isel(array.indices_max(dim=("x", "z")))
+        >>> array.isel(array.argmax(dim=["x", "z"]))
         <xarray.DataArray (y: 3)>
         array([3, 5, 3])
         Dimensions without coordinates: y
         """
-        return self._unravel_argminmax("argmax", dim, keep_attrs, skipna)
+        return self._unravel_argminmax("_argmax_base", dim, axis, keep_attrs, skipna)
 
 
 # priority most be higher than Variable to properly work with binary ufuncs

xarray/core/duck_array_ops.py

@@ -319,8 +319,8 @@ def f(values, axis=None, skipna=None, **kwargs):
 
 # Attributes `numeric_only`, `available_min_count` is used for docs.
 # See ops.inject_reduce_methods
-argmax = _create_nan_agg_method("argmax", coerce_strings=True)
-argmin = _create_nan_agg_method("argmin", coerce_strings=True)
+_argmax_base = _create_nan_agg_method("argmax", coerce_strings=True)
+_argmin_base = _create_nan_agg_method("argmin", coerce_strings=True)
 max = _create_nan_agg_method("max", coerce_strings=True)
 min = _create_nan_agg_method("min", coerce_strings=True)
 sum = _create_nan_agg_method("sum")

xarray/core/ops.py

@@ -47,8 +47,8 @@
 # methods which remove an axis
 REDUCE_METHODS = ["all", "any"]
 NAN_REDUCE_METHODS = [
-    "argmax",
-    "argmin",
+    "_argmax_base",
+    "_argmin_base",
     "max",
     "min",
     "mean",

xarray/core/rolling.py

@@ -130,8 +130,8 @@ def method(self, **kwargs):
         method.__doc__ = _ROLLING_REDUCE_DOCSTRING_TEMPLATE.format(name=name)
         return method
 
-    argmax = _reduce_method("argmax")
-    argmin = _reduce_method("argmin")
+    _argmax_base = _reduce_method("_argmax_base")
+    _argmin_base = _reduce_method("_argmin_base")
     max = _reduce_method("max")
     min = _reduce_method("min")
     mean = _reduce_method("mean")