implement idxmax and idxmin

doc/api.rst

@@ -180,6 +180,8 @@ Computation
 :py:attr:`~Dataset.any`
 :py:attr:`~Dataset.argmax`
 :py:attr:`~Dataset.argmin`
+:py:attr:`~Dataset.idxmax`
+:py:attr:`~Dataset.idxmin`
 :py:attr:`~Dataset.max`
 :py:attr:`~Dataset.mean`
 :py:attr:`~Dataset.median`
@@ -362,6 +364,8 @@ Computation
 :py:attr:`~DataArray.any`
 :py:attr:`~DataArray.argmax`
 :py:attr:`~DataArray.argmin`
+:py:attr:`~DataArray.idxmax`
+:py:attr:`~DataArray.idxmin`
 :py:attr:`~DataArray.max`
 :py:attr:`~DataArray.mean`
 :py:attr:`~DataArray.median`

doc/whats-new.rst

@@ -38,6 +38,9 @@ New Features
 - Limited the length of array items with long string reprs to a
   reasonable width (:pull:`3900`)
   By `Maximilian Roos <https://github.com/max-sixty>`_
+- Implement :py:meth:`DataArray.idxmax`, :py:meth:`DataArray.idxmin`,
+  :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`.  (:issue:`60`, :pull:`3871`)
+  By `Todd Jennings <https://github.com/toddrjen>`_
 
 
 Bug fixes

xarray/core/computation.py

@@ -23,9 +23,10 @@
 
 import numpy as np
 
-from . import duck_array_ops, utils
+from . import dtypes, duck_array_ops, utils
 from .alignment import deep_align
 from .merge import merge_coordinates_without_align
+from .nanops import dask_array
 from .options import OPTIONS
 from .pycompat import dask_array_type
 from .utils import is_dict_like
@@ -1338,3 +1339,66 @@ def polyval(coord, coeffs, degree_dim="degree"):
         coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]},
     )
     return (lhs * coeffs).sum(degree_dim)
+
+
+def _calc_idxminmax(
+    *,
+    array,
+    func: Callable,
+    dim: Hashable = None,
+    skipna: bool = None,
+    fill_value: Any = dtypes.NA,
+    keep_attrs: bool = None,
+):
+    """Apply common operations for idxmin and idxmax."""
+    # This function doesn't make sense for scalars so don't try
+    if not array.ndim:
+        raise ValueError("This function does not apply for scalars")
+
+    if dim is not None:
+        pass  # Use the dim if available
+    elif array.ndim == 1:
+        # it is okay to guess the dim if there is only 1
+        dim = array.dims[0]
+    else:
+        # The dim is not specified and ambiguous.  Don't guess.
+        raise ValueError("Must supply 'dim' argument for multidimensional arrays")
+
+    if dim not in array.dims:
+        raise KeyError(f'Dimension "{dim}" not in dimension')
+    if dim not in array.coords:
+        raise KeyError(f'Dimension "{dim}" does not have coordinates')
+
+    # These are dtypes with NaN values argmin and argmax can handle
+    na_dtypes = "cfO"
+
+    if skipna or (skipna is None and array.dtype.kind in na_dtypes):
+        # Need to skip NaN values since argmin and argmax can't handle them
+        allna = array.isnull().all(dim)
+        array = array.where(~allna, 0)
+
+    # This will run argmin or argmax.
+    indx = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna)
+
+    # Get the coordinate we want.
+    coordarray = array[dim]
+
+    # Handle dask arrays.
+    if isinstance(array, dask_array_type):
+        res = dask_array.map_blocks(coordarray, indx, dtype=indx.dtype)
+    else:
+        res = coordarray[
+            indx,
+        ]
+
+    if skipna or (skipna is None and array.dtype.kind in na_dtypes):
+        # Put the NaN values back in after removing them
+        res = res.where(~allna, fill_value)
+
+    # The dim is gone but we need to remove the corresponding coordinate.
+    del res.coords[dim]
+
+    # Copy attributes from argmin/argmax, if any
+    res.attrs = indx.attrs
+
+    return res

xarray/core/dataarray.py

@@ -3508,6 +3508,194 @@ def pad(
         )
         return self._from_temp_dataset(ds)
 
+    def idxmin(
+        self,
+        dim: Hashable = None,
+        skipna: bool = None,
+        fill_value: Any = dtypes.NA,
+        keep_attrs: bool = None,
+    ) -> "DataArray":
+        """Return the coordinate label of the minimum value along a dimension.
+
+        Returns a new `DataArray` named after the dimension with the values of
+        the coordinate labels along that dimension corresponding to minimum
+        values along that dimension.
+
+        In comparison to :py:meth:`~DataArray.argmin`, this returns the
+        coordinate label while :py:meth:`~DataArray.argmin` returns the index.
+
+        Parameters
+        ----------
+        dim : str, optional
+            Dimension over which to apply `idxmin`.  This is optional for 1D
+            arrays, but required for arrays with 2 or more dimensions.
+        skipna : bool or None, default None
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for ``float``, ``complex``, and ``object``
+            dtypes; other dtypes either do not have a sentinel missing value
+            (``int``) or ``skipna=True`` has not been implemented
+            (``datetime64`` or ``timedelta64``).
+        fill_value : Any, default NaN
+            Value to be filled in case all of the values along a dimension are
+            null.  By default this is NaN.  The fill value and result are
+            automatically converted to a compatible dtype if possible.
+            Ignored if ``skipna`` is False.
+        keep_attrs : bool, default False
+            If True, the attributes (``attrs``) will be copied from the
+            original object to the new one.  If False (default), the new object
+            will be returned without attributes.
+
+        Returns
+        -------
+        reduced : DataArray
+            New `DataArray` object with `idxmin` applied to its data and the
+            indicated dimension removed.
+
+        See also
+        --------
+        Dataset.idxmin, DataArray.idxmax, DataArray.min, DataArray.argmin
+
+        Examples
+        --------
+
+        >>> array = xr.DataArray([0, 2, 1, 0, -2], dims="x",
+        ...                      coords={"x": ['a', 'b', 'c', 'd', 'e']})
+        >>> array.min()
+        <xarray.DataArray ()>
+        array(-2)
+        >>> array.argmin()
+        <xarray.DataArray ()>
+        array(4)
+        >>> array.idxmin()
+        <xarray.DataArray 'x' ()>
+        array('e', dtype='<U1')
+
+        >>> array = xr.DataArray([[2.0, 1.0, 2.0, 0.0, -2.0],
+        ...                       [-4.0, np.NaN, 2.0, np.NaN, -2.0],
+        ...                       [np.NaN, np.NaN, 1., np.NaN, np.NaN]],
+        ...                      dims=["y", "x"],
+        ...                      coords={"y": [-1, 0, 1],
+        ...                              "x": np.arange(5.)**2}
+        ...                      )
+        >>> array.min(dim="x")
+        <xarray.DataArray (y: 3)>
+        array([-2., -4.,  1.])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        >>> array.argmin(dim="x")
+        <xarray.DataArray (y: 3)>
+        array([4, 0, 2])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        >>> array.idxmin(dim="x")
+        <xarray.DataArray 'x' (y: 3)>
+        array([16.,  0.,  4.])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        """
+        return computation._calc_idxminmax(
+            array=self,
+            func=lambda x, *args, **kwargs: x.argmin(*args, **kwargs),
+            dim=dim,
+            skipna=skipna,
+            fill_value=fill_value,
+            keep_attrs=keep_attrs,
+        )
+
+    def idxmax(
+        self,
+        dim: Hashable = None,
+        skipna: bool = None,
+        fill_value: Any = dtypes.NA,
+        keep_attrs: bool = None,
+    ) -> "DataArray":
+        """Return the coordinate label of the maximum value along a dimension.
+
+        Returns a new `DataArray` named after the dimension with the values of
+        the coordinate labels along that dimension corresponding to maximum
+        values along that dimension.
+
+        In comparison to :py:meth:`~DataArray.argmax`, this returns the
+        coordinate label while :py:meth:`~DataArray.argmax` returns the index.
+
+        Parameters
+        ----------
+        dim : str, optional
+            Dimension over which to apply `idxmax`.  This is optional for 1D
+            arrays, but required for arrays with 2 or more dimensions.
+        skipna : bool or None, default None
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for ``float``, ``complex``, and ``object``
+            dtypes; other dtypes either do not have a sentinel missing value
+            (``int``) or ``skipna=True`` has not been implemented
+            (``datetime64`` or ``timedelta64``).
+        fill_value : Any, default NaN
+            Value to be filled in case all of the values along a dimension are
+            null.  By default this is NaN.  The fill value and result are
+            automatically converted to a compatible dtype if possible.
+            Ignored if ``skipna`` is False.
+        keep_attrs : bool, default False
+            If True, the attributes (``attrs``) will be copied from the
+            original object to the new one.  If False (default), the new object
+            will be returned without attributes.
+
+        Returns
+        -------
+        reduced : DataArray
+            New `DataArray` object with `idxmax` applied to its data and the
+            indicated dimension removed.
+
+        See also
+        --------
+        Dataset.idxmax, DataArray.idxmin, DataArray.max, DataArray.argmax
+
+        Examples
+        --------
+
+        >>> array = xr.DataArray([0, 2, 1, 0, -2], dims="x",
+        ...                      coords={"x": ['a', 'b', 'c', 'd', 'e']})
+        >>> array.max()
+        <xarray.DataArray ()>
+        array(2)
+        >>> array.argmax()
+        <xarray.DataArray ()>
+        array(1)
+        >>> array.idxmax()
+        <xarray.DataArray 'x' ()>
+        array('b', dtype='<U1')
+
+        >>> array = xr.DataArray([[2.0, 1.0, 2.0, 0.0, -2.0],
+        ...                       [-4.0, np.NaN, 2.0, np.NaN, -2.0],
+        ...                       [np.NaN, np.NaN, 1., np.NaN, np.NaN]],
+        ...                      dims=["y", "x"],
+        ...                      coords={"y": [-1, 0, 1],
+        ...                              "x": np.arange(5.)**2}
+        ...                      )
+        >>> array.max(dim="x")
+        <xarray.DataArray (y: 3)>
+        array([2., 2., 1.])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        >>> array.argmax(dim="x")
+        <xarray.DataArray (y: 3)>
+        array([0, 2, 2])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        >>> array.idxmax(dim="x")
+        <xarray.DataArray 'x' (y: 3)>
+        array([0., 4., 4.])
+        Coordinates:
+          * y        (y) int64 -1 0 1
+        """
+        return computation._calc_idxminmax(
+            array=self,
+            func=lambda x, *args, **kwargs: x.argmax(*args, **kwargs),
+            dim=dim,
+            skipna=skipna,
+            fill_value=fill_value,
+            keep_attrs=keep_attrs,
+        )
+
     # this needs to be at the end, or mypy will confuse with `str`
     # https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names
     str = property(StringAccessor)