| 5840 | assert_identical(result3, expected2) |
| 5841 | |
| 5842 | def test_argmax( |
| 5843 | self, |
| 5844 | x: np.ndarray, |
| 5845 | minindex: list[int | float], |
| 5846 | maxindex: list[int | float], |
| 5847 | nanindex: list[int | None], |
| 5848 | ) -> None: |
| 5849 | ar = xr.DataArray( |
| 5850 | x, |
| 5851 | dims=["y", "x"], |
| 5852 | coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, |
| 5853 | attrs=self.attrs, |
| 5854 | ) |
| 5855 | indarr_np = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) |
| 5856 | indarr = xr.DataArray(indarr_np, dims=ar.dims, coords=ar.coords) |
| 5857 | |
| 5858 | if np.isnan(maxindex).any(): |
| 5859 | with pytest.raises(ValueError): |
| 5860 | ar.argmax(dim="x") |
| 5861 | return |
| 5862 | |
| 5863 | expected0list = [ |
| 5864 | indarr.isel(y=yi).isel(x=indi, drop=True) |
| 5865 | for yi, indi in enumerate(maxindex) |
| 5866 | ] |
| 5867 | expected0 = xr.concat(expected0list, dim="y") |
| 5868 | expected0.attrs = self.attrs # Default keeps attrs for reduction operations |
| 5869 | |
| 5870 | result0 = ar.argmax(dim="x") |
| 5871 | assert_identical(result0, expected0) |
| 5872 | |
| 5873 | result1 = ar.argmax(axis=1) |
| 5874 | assert_identical(result1, expected0) |
| 5875 | |
| 5876 | result2 = ar.argmax(dim="x", keep_attrs=True) |
| 5877 | expected1 = expected0.copy() |
| 5878 | expected1.attrs = self.attrs |
| 5879 | assert_identical(result2, expected1) |
| 5880 | |
| 5881 | maxindex = [ |
| 5882 | x if y is None or ar.dtype.kind == "O" else y |
| 5883 | for x, y in zip(maxindex, nanindex, strict=True) |
| 5884 | ] |
| 5885 | expected2list = [ |
| 5886 | indarr.isel(y=yi).isel(x=indi, drop=True) |
| 5887 | for yi, indi in enumerate(maxindex) |
| 5888 | ] |
| 5889 | expected2 = xr.concat(expected2list, dim="y") |
| 5890 | expected2.attrs = self.attrs # Default keeps attrs for reduction operations |
| 5891 | |
| 5892 | result3 = ar.argmax(dim="x", skipna=False) |
| 5893 | |
| 5894 | assert_identical(result3, expected2) |
| 5895 | |
| 5896 | @pytest.mark.parametrize( |
| 5897 | "use_dask", [pytest.param(True, id="dask"), pytest.param(False, id="nodask")] |