Skip to content

Commit 7d65502

Browse files
toddrjentoddrjen
committed
implement idxmax and idxmin
1 parent 70e628d commit 7d65502

File tree

7 files changed

+1194
-1
lines changed

7 files changed

+1194
-1
lines changed

‎doc/api.rst

+4
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,8 @@ Computation
180180
:py:attr:`~Dataset.any`
181181
:py:attr:`~Dataset.argmax`
182182
:py:attr:`~Dataset.argmin`
183+
:py:attr:`~Dataset.idxmax`
184+
:py:attr:`~Dataset.idxmin`
183185
:py:attr:`~Dataset.max`
184186
:py:attr:`~Dataset.mean`
185187
:py:attr:`~Dataset.median`
@@ -362,6 +364,8 @@ Computation
362364
:py:attr:`~DataArray.any`
363365
:py:attr:`~DataArray.argmax`
364366
:py:attr:`~DataArray.argmin`
367+
:py:attr:`~DataArray.idxmax`
368+
:py:attr:`~DataArray.idxmin`
365369
:py:attr:`~DataArray.max`
366370
:py:attr:`~DataArray.mean`
367371
:py:attr:`~DataArray.median`

‎doc/whats-new.rst

+3
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ New Features
3838
- Limited the length of array items with long string reprs to a
3939
reasonable width (:pull:`3900`)
4040
By `Maximilian Roos <https://github.com/max-sixty>`_
41+
- Implement :py:meth:`DataArray.idxmax`, :py:meth:`DataArray.idxmin`,
42+
:py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:issue:`60`, :pull:`3871`)
43+
By `Todd Jennings <https://github.com/toddrjen>`_
4144

4245

4346
Bug fixes

‎xarray/core/computation.py

+65-1
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,10 @@
2323

2424
import numpy as np
2525

26-
from . import duck_array_ops, utils
26+
from . import dtypes, duck_array_ops, utils
2727
from .alignment import deep_align
2828
from .merge import merge_coordinates_without_align
29+
from .nanops import dask_array
2930
from .options import OPTIONS
3031
from .pycompat import dask_array_type
3132
from .utils import is_dict_like
@@ -1338,3 +1339,66 @@ def polyval(coord, coeffs, degree_dim="degree"):
13381339
coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]},
13391340
)
13401341
return (lhs * coeffs).sum(degree_dim)
1342+
1343+
1344+
def _calc_idxminmax(
1345+
*,
1346+
array,
1347+
func: Callable,
1348+
dim: Hashable = None,
1349+
skipna: bool = None,
1350+
fill_value: Any = dtypes.NA,
1351+
keep_attrs: bool = None,
1352+
):
1353+
"""Apply common operations for idxmin and idxmax."""
1354+
# This function doesn't make sense for scalars so don't try
1355+
if not array.ndim:
1356+
raise ValueError("This function does not apply for scalars")
1357+
1358+
if dim is not None:
1359+
pass # Use the dim if available
1360+
elif array.ndim == 1:
1361+
# it is okay to guess the dim if there is only 1
1362+
dim = array.dims[0]
1363+
else:
1364+
# The dim is not specified and ambiguous. Don't guess.
1365+
raise ValueError("Must supply 'dim' argument for multidimensional arrays")
1366+
1367+
if dim not in array.dims:
1368+
raise KeyError(f'Dimension "{dim}" not in dimension')
1369+
if dim not in array.coords:
1370+
raise KeyError(f'Dimension "{dim}" does not have coordinates')
1371+
1372+
# These are dtypes with NaN values argmin and argmax can handle
1373+
na_dtypes = "cfO"
1374+
1375+
if skipna or (skipna is None and array.dtype.kind in na_dtypes):
1376+
# Need to skip NaN values since argmin and argmax can't handle them
1377+
allna = array.isnull().all(dim)
1378+
array = array.where(~allna, 0)
1379+
1380+
# This will run argmin or argmax.
1381+
indx = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna)
1382+
1383+
# Get the coordinate we want.
1384+
coordarray = array[dim]
1385+
1386+
# Handle dask arrays.
1387+
if isinstance(array, dask_array_type):
1388+
res = dask_array.map_blocks(coordarray, indx, dtype=indx.dtype)
1389+
else:
1390+
res = coordarray[
1391+
indx,
1392+
]
1393+
1394+
if skipna or (skipna is None and array.dtype.kind in na_dtypes):
1395+
# Put the NaN values back in after removing them
1396+
res = res.where(~allna, fill_value)
1397+
1398+
# The dim is gone but we need to remove the corresponding coordinate.
1399+
del res.coords[dim]
1400+
1401+
# Copy attributes from argmin/argmax, if any
1402+
res.attrs = indx.attrs
1403+
1404+
return res

‎xarray/core/dataarray.py

+185
Original file line numberDiff line numberDiff line change
@@ -3508,6 +3508,191 @@ def pad(
35083508
)
35093509
return self._from_temp_dataset(ds)
35103510

3511+
def idxmin(
3512+
self,
3513+
dim: Hashable = None,
3514+
skipna: bool = None,
3515+
fill_value: Any = dtypes.NA,
3516+
keep_attrs: bool = None,
3517+
) -> "DataArray":
3518+
"""Return the coordinate of the minimum value along a dimension.
3519+
3520+
Returns a new DataArray named after the dimension with the values of
3521+
the coordinate along that dimension corresponding to minimum value
3522+
along that dimension.
3523+
3524+
In comparison to `argmin`, this returns the coordinate label while
3525+
`argmin` the index.
3526+
3527+
Parameters
3528+
----------
3529+
dim : str, optional
3530+
Dimension over which to apply `idxmin`. This is optional for 1D
3531+
arrays, but required for arrays with 2 or more dimensions.
3532+
skipna : bool or None, default None
3533+
If True, skip missing values (as marked by NaN). By default, only
3534+
skips missing values for float dtypes; other dtypes either do not
3535+
have a sentinel missing value (int) or skipna=True has not been
3536+
implemented (object, datetime64 or timedelta64).
3537+
fill_value : Any, default NaN
3538+
Value to be filled in case all of the values along a dimension are
3539+
null. By default this is NaN. The fill value and result are
3540+
automatically converted to a compatible dtype if possible.
3541+
Ignored if `skipna` is False.
3542+
keep_attrs : bool, default False
3543+
If True, the attributes (`attrs`) will be copied from the original
3544+
object to the new one. If False (default), the new object will be
3545+
returned without attributes.
3546+
3547+
Returns
3548+
-------
3549+
reduced : DataArray
3550+
New DataArray object with `idxmin` applied to its data and the
3551+
indicated dimension removed.
3552+
3553+
Examples
3554+
--------
3555+
3556+
>>> array = xr.DataArray([0, 2, 1, 0, -2], dims="x",
3557+
... coords={"x": ['a', 'b', 'c', 'd', 'e']})
3558+
>>> array.min()
3559+
<xarray.DataArray ()>
3560+
array(-2)
3561+
>>> array.argmin()
3562+
<xarray.DataArray ()>
3563+
array(4)
3564+
>>> array.idxmin()
3565+
<xarray.DataArray 'x' ()>
3566+
array('e', dtype='<U1')
3567+
3568+
>>> array = xr.DataArray([[2.0, 1.0, 2.0, 0.0, -2.0],
3569+
... [-4.0, np.NaN, 2.0, np.NaN, -2.0],
3570+
... [np.NaN, np.NaN, 1., np.NaN, np.NaN]],
3571+
... dims=["y", "x"],
3572+
... coords={"y": [-1, 0, 1],
3573+
... "x": np.arange(5.)**2}
3574+
... )
3575+
>>> array.min(dim="x")
3576+
<xarray.DataArray (y: 3)>
3577+
array([-2., -4., 1.])
3578+
Coordinates:
3579+
* y (y) int64 -1 0 1
3580+
>>> array.argmin(dim="x")
3581+
<xarray.DataArray (y: 3)>
3582+
array([4, 0, 2])
3583+
Coordinates:
3584+
* y (y) int64 -1 0 1
3585+
>>> array.idxmin(dim="x")
3586+
<xarray.DataArray 'x' (y: 3)>
3587+
array([16., 0., 4.])
3588+
Coordinates:
3589+
* y (y) int64 -1 0 1
3590+
3591+
See also
3592+
--------
3593+
Dataset.idxmin, DataArray.idxmax, DataArray.min, DataArray.argmin
3594+
"""
3595+
return computation._calc_idxminmax(
3596+
array=self,
3597+
func=lambda x, *args, **kwargs: x.argmin(*args, **kwargs),
3598+
dim=dim,
3599+
skipna=skipna,
3600+
fill_value=fill_value,
3601+
keep_attrs=keep_attrs,
3602+
)
3603+
3604+
def idxmax(
3605+
self,
3606+
dim: Hashable = None,
3607+
skipna: bool = None,
3608+
fill_value: Any = dtypes.NA,
3609+
keep_attrs: bool = None,
3610+
) -> "DataArray":
3611+
"""Return the coordinate of the maximum value along a dimension.
3612+
3613+
Returns a new DataArray named after the dimension with the values of
3614+
the coordinate along that dimension corresponding to maximum value
3615+
along that dimension.
3616+
3617+
In comparison to `argmax`, this returns the coordinate label while
3618+
`argmax` returns the index.
3619+
3620+
Parameters
3621+
----------
3622+
dim : str, optional
3623+
Dimension over which to apply `idxmax`. This is optional for 1D
3624+
arrays, but required for arrays with 2 or more dimensions.
3625+
skipna : bool or None, default None
3626+
If True, skip missing values (as marked by NaN). By default, only
3627+
skips missing values for float dtypes; other dtypes either do not
3628+
have a sentinel missing value (int) or skipna=True has not been
3629+
implemented (object, datetime64 or timedelta64).
3630+
fill_value : Any, default NaN
3631+
Value to be filled in case all of the values along a dimension are
3632+
null. By default this is NaN. The fill value and result are
3633+
automatically converted to a compatible dtype if possible.
3634+
Ignored if `skipna` is False.
3635+
keep_attrs : bool, default False
3636+
If True, the attributes (`attrs`) will be copied from the original
3637+
object to the new one. If False (default), the new object will be
3638+
returned without attributes.
3639+
Returns
3640+
-------
3641+
reduced : DataArray
3642+
New DataArray object with `idxmax` applied to its data and the
3643+
indicated dimension removed.
3644+
3645+
Examples
3646+
--------
3647+
3648+
>>> array = xr.DataArray([0, 2, 1, 0, -2], dims="x",
3649+
... coords={"x": ['a', 'b', 'c', 'd', 'e']})
3650+
>>> array.max()
3651+
<xarray.DataArray ()>
3652+
array(2)
3653+
>>> array.argmax()
3654+
<xarray.DataArray ()>
3655+
array(1)
3656+
>>> array.idxmax()
3657+
<xarray.DataArray 'x' ()>
3658+
array('b', dtype='<U1')
3659+
3660+
>>> array = xr.DataArray([[2.0, 1.0, 2.0, 0.0, -2.0],
3661+
... [-4.0, np.NaN, 2.0, np.NaN, -2.0],
3662+
... [np.NaN, np.NaN, 1., np.NaN, np.NaN]],
3663+
... dims=["y", "x"],
3664+
... coords={"y": [-1, 0, 1],
3665+
... "x": np.arange(5.)**2}
3666+
... )
3667+
>>> array.max(dim="x")
3668+
<xarray.DataArray (y: 3)>
3669+
array([2., 2., 1.])
3670+
Coordinates:
3671+
* y (y) int64 -1 0 1
3672+
>>> array.argmax(dim="x")
3673+
<xarray.DataArray (y: 3)>
3674+
array([0, 2, 2])
3675+
Coordinates:
3676+
* y (y) int64 -1 0 1
3677+
>>> array.idxmax(dim="x")
3678+
<xarray.DataArray 'x' (y: 3)>
3679+
array([0., 4., 4.])
3680+
Coordinates:
3681+
* y (y) int64 -1 0 1
3682+
3683+
See also
3684+
--------
3685+
Dataset.idxmax, DataArray.idxmin, DataArray.max, DataArray.argmax
3686+
"""
3687+
return computation._calc_idxminmax(
3688+
array=self,
3689+
func=lambda x, *args, **kwargs: x.argmax(*args, **kwargs),
3690+
dim=dim,
3691+
skipna=skipna,
3692+
fill_value=fill_value,
3693+
keep_attrs=keep_attrs,
3694+
)
3695+
35113696
# this needs to be at the end, or mypy will confuse with `str`
35123697
# https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names
35133698
str = property(StringAccessor)

0 commit comments

Comments
 (0)