| 5413 | |
| 5414 | @pytest.mark.parametrize("use_dask", [True, False]) |
| 5415 | def test_idxmax( |
| 5416 | self, |
| 5417 | x: np.ndarray, |
| 5418 | minindex: int | float, |
| 5419 | maxindex: int | float, |
| 5420 | nanindex: int | None, |
| 5421 | use_dask: bool, |
| 5422 | ) -> None: |
| 5423 | if use_dask and not has_dask: |
| 5424 | pytest.skip("requires dask") |
| 5425 | if use_dask and x.dtype.kind == "M": |
| 5426 | pytest.xfail("dask operation 'argmax' breaks when dtype is datetime64 (M)") |
| 5427 | ar0_raw = xr.DataArray( |
| 5428 | x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs |
| 5429 | ) |
| 5430 | |
| 5431 | if use_dask: |
| 5432 | ar0 = ar0_raw.chunk({}) |
| 5433 | else: |
| 5434 | ar0 = ar0_raw |
| 5435 | |
| 5436 | with pytest.raises( |
| 5437 | KeyError, |
| 5438 | match=r"'spam' not found in array dimensions", |
| 5439 | ): |
| 5440 | ar0.idxmax(dim="spam") |
| 5441 | |
| 5442 | # Scalar Dataarray |
| 5443 | with pytest.raises(ValueError): |
| 5444 | xr.DataArray(5).idxmax() |
| 5445 | |
| 5446 | coordarr0 = xr.DataArray(ar0.coords["x"].data, dims=["x"]) |
| 5447 | coordarr1 = coordarr0.copy() |
| 5448 | |
| 5449 | hasna = np.isnan(maxindex) |
| 5450 | if np.isnan(maxindex): |
| 5451 | maxindex = 0 |
| 5452 | |
| 5453 | if hasna: |
| 5454 | coordarr1[...] = 1 |
| 5455 | fill_value_0 = np.nan |
| 5456 | else: |
| 5457 | fill_value_0 = 1 |
| 5458 | |
| 5459 | expected0 = ( |
| 5460 | (coordarr1 * fill_value_0).isel(x=maxindex, drop=True).astype("float") |
| 5461 | ) |
| 5462 | expected0.name = "x" |
| 5463 | expected0.attrs = self.attrs # Default keeps attrs for reduction operations |
| 5464 | |
| 5465 | # Default fill value (NaN) |
| 5466 | result0 = ar0.idxmax() |
| 5467 | assert_identical(result0, expected0) |
| 5468 | |
| 5469 | # Manually specify NaN fill_value |
| 5470 | result1 = ar0.idxmax(fill_value=np.nan) |
| 5471 | assert_identical(result1, expected0) |
| 5472 | |