Function _calc_idxminmax

xarray/computation/computation.py:960–1023

Apply common operations for idxmin and idxmax.

(
    *,
    array,
    func: Callable,
    dim: Hashable | None = None,
    skipna: bool | None = None,
    fill_value: Any = dtypes.NA,
    keep_attrs: bool | None = None,
)

def _calc_idxminmax(
    *,
    array,
    func: Callable,
    dim: Hashable | None = None,
    skipna: bool | None = None,
    fill_value: Any = dtypes.NA,
    keep_attrs: bool | None = 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!r} not found in array dimensions {array.dims!r}"
        )
    if dim not in array.coords:
        raise KeyError(
            f"Dimension {dim!r} is not one of the coordinates {tuple(array.coords.keys())}"
        )

    # 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.
    index = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna)

    # Handle chunked arrays (e.g. dask).
    coord = array[dim]._variable.to_base_variable()
    if is_chunked_array(array.data):
        chunkmanager = get_chunked_array_type(array.data)
        coord_array = chunkmanager.from_array(
            array[dim].data, chunks=((array.sizes[dim],),)
        )
        coord = coord.copy(data=coord_array)
    else:
        coord = coord.copy(data=to_like_array(array[dim].data, array.data))

    res = index._replace(coord[(index.variable,)]).rename(dim)

    if skipna or (skipna is None and array.dtype.kind in na_dtypes):
        # Put the NaN values back in after removing them