Method test_idxmax

xarray/tests/test_dataarray.py:6044–6185

(
        self,
        x: np.ndarray,
        minindex: list[int | float],
        maxindex: list[int | float],
        nanindex: list[int | None],
        use_dask: bool,
    )

        "use_dask", [pytest.param(True, id="dask"), pytest.param(False, id="nodask")]
    )
    def test_idxmax(
        self,
        x: np.ndarray,
        minindex: list[int | float],
        maxindex: list[int | float],
        nanindex: list[int | None],
        use_dask: bool,
    ) -> None:
        if use_dask and not has_dask:
            pytest.skip("requires dask")
        if use_dask and x.dtype.kind == "M":
            pytest.xfail("dask operation 'argmax' breaks when dtype is datetime64 (M)")

        if x.dtype.kind == "O":
            # TODO: nanops._nan_argminmax_object computes once to check for all-NaN slices.
            max_computes = 1
        else:
            max_computes = 0

        ar0_raw = xr.DataArray(
            x,
            dims=["y", "x"],
            coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])},
            attrs=self.attrs,
        )

        if use_dask:
            ar0 = ar0_raw.chunk({})
        else:
            ar0 = ar0_raw

        # No dimension specified
        with pytest.raises(ValueError):
            ar0.idxmax()

        # dim doesn't exist
        with pytest.raises(KeyError):
            ar0.idxmax(dim="Y")

        ar1 = ar0.copy()
        del ar1.coords["y"]
        with pytest.raises(KeyError):
            ar1.idxmax(dim="y")

        coordarr0 = xr.DataArray(
            np.tile(ar0.coords["x"], [x.shape[0], 1]), dims=ar0.dims, coords=ar0.coords
        )

        hasna = [np.isnan(x) for x in maxindex]
        coordarr1 = coordarr0.copy()
        coordarr1[hasna, :] = 1
        maxindex0 = [x if not np.isnan(x) else 0 for x in maxindex]

        nan_mult_0 = np.array([np.nan if x else 1 for x in hasna])[:, None]
        expected0list = [
            (coordarr1 * nan_mult_0).isel(y=yi).isel(x=indi, drop=True)
            for yi, indi in enumerate(maxindex0)
        ]