| 6260 | "ignore:Behaviour of argmin/argmax with neither dim nor :FutureWarning" |
| 6261 | ) |
| 6262 | def test_argmax_dim( |
| 6263 | self, |
| 6264 | x: np.ndarray, |
| 6265 | minindex: list[int | float], |
| 6266 | maxindex: list[int | float], |
| 6267 | nanindex: list[int | None], |
| 6268 | ) -> None: |
| 6269 | ar = xr.DataArray( |
| 6270 | x, |
| 6271 | dims=["y", "x"], |
| 6272 | coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, |
| 6273 | attrs=self.attrs, |
| 6274 | ) |
| 6275 | indarrnp = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) |
| 6276 | indarr = xr.DataArray(indarrnp, dims=ar.dims, coords=ar.coords) |
| 6277 | |
| 6278 | if np.isnan(maxindex).any(): |
| 6279 | with pytest.raises(ValueError): |
| 6280 | ar.argmax(dim="x") |
| 6281 | return |
| 6282 | |
| 6283 | expected0list = [ |
| 6284 | indarr.isel(y=yi).isel(x=indi, drop=True) |
| 6285 | for yi, indi in enumerate(maxindex) |
| 6286 | ] |
| 6287 | expected0 = {"x": xr.concat(expected0list, dim="y")} |
| 6288 | expected0[ |
| 6289 | "x" |
| 6290 | ].attrs = self.attrs # Default keeps attrs for reduction operations |
| 6291 | |
| 6292 | result0 = ar.argmax(dim=["x"]) |
| 6293 | for key in expected0: |
| 6294 | assert_identical(result0[key], expected0[key]) |
| 6295 | |
| 6296 | result1 = ar.argmax(dim=["x"], keep_attrs=True) |
| 6297 | expected1 = deepcopy(expected0) |
| 6298 | expected1["x"].attrs = self.attrs |
| 6299 | for key in expected1: |
| 6300 | assert_identical(result1[key], expected1[key]) |
| 6301 | |
| 6302 | maxindex = [ |
| 6303 | x if y is None or ar.dtype.kind == "O" else y |
| 6304 | for x, y in zip(maxindex, nanindex, strict=True) |
| 6305 | ] |
| 6306 | expected2list = [ |
| 6307 | indarr.isel(y=yi).isel(x=indi, drop=True) |
| 6308 | for yi, indi in enumerate(maxindex) |
| 6309 | ] |
| 6310 | expected2 = {"x": xr.concat(expected2list, dim="y")} |
| 6311 | expected2[ |
| 6312 | "x" |
| 6313 | ].attrs = self.attrs # Default keeps attrs for reduction operations |
| 6314 | |
| 6315 | result2 = ar.argmax(dim=["x"], skipna=False) |
| 6316 | |
| 6317 | for key in expected2: |
| 6318 | assert_identical(result2[key], expected2[key]) |
| 6319 | |