From 8510a3161c28ee1287aa21e3428c74aeb2ae7509 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sun, 31 Aug 2025 12:33:16 -0400 Subject: [PATCH 1/7] allow multiple dims --- xarray/computation/computation.py | 65 ++++++++++++++----------------- 1 file changed, 30 insertions(+), 35 deletions(-) diff --git a/xarray/computation/computation.py b/xarray/computation/computation.py index 14b1ae6e240..2d97bc10458 100644 --- a/xarray/computation/computation.py +++ b/xarray/computation/computation.py @@ -9,10 +9,7 @@ import functools from collections import Counter -from collections.abc import ( - Callable, - Hashable, -) +from collections.abc import Callable, Hashable from typing import TYPE_CHECKING, Any, Literal, cast, overload import numpy as np @@ -23,10 +20,7 @@ from xarray.core.duck_array_ops import datetime_to_numeric from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.types import Dims, T_DataArray -from xarray.core.utils import ( - is_scalar, - parse_dims_as_set, -) +from xarray.core.utils import is_scalar, parse_dims_as_set from xarray.core.variable import Variable from xarray.namedarray.parallelcompat import get_chunked_array_type from xarray.namedarray.pycompat import is_chunked_array @@ -912,14 +906,15 @@ def _calc_idxminmax( # 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())}" - ) + for _dim in dim: + 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" @@ -932,24 +927,24 @@ def _calc_idxminmax( # This will run argmin or argmax. indx = 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 = indx._replace(coord[(indx.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 - res = res.where(~allna, fill_value) - - # Copy attributes from argmin/argmax, if any - res.attrs = indx.attrs + res = {} + for _dim, _da_idx in zip(dim, indx.values(), strict=False): + # 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 = _da_idx._replace(coord[(_da_idx.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 + _res = _res.where(~allna, fill_value) + _res.attrs = _da_idx.attrs + res[_dim] = _res return res From aaa66a233294e2e9c53761437c0d717e75fd13cc Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sun, 31 Aug 2025 14:53:32 -0400 Subject: [PATCH 2/7] adapt code for case dim==1 --- xarray/computation/computation.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/xarray/computation/computation.py b/xarray/computation/computation.py index 2d97bc10458..074254a0d3c 100644 --- a/xarray/computation/computation.py +++ b/xarray/computation/computation.py @@ -906,7 +906,9 @@ def _calc_idxminmax( # The dim is not specified and ambiguous. Don't guess. raise ValueError("Must supply 'dim' argument for multidimensional arrays") - for _dim in dim: + dims = [dim] if isinstance(dim, str) else list(dim) + + for _dim in dims: if _dim not in array.dims: raise KeyError( f"Dimension {_dim!r} not found in array dimensions {array.dims!r}" @@ -926,9 +928,12 @@ def _calc_idxminmax( # This will run argmin or argmax. indx = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna) + # Force dictionary format in case of single dim so that we can iterate over it in for loop below + if len(dims) == 1: + indx = {dims[0]: indx} res = {} - for _dim, _da_idx in zip(dim, indx.values(), strict=False): + for _dim, _da_idx in zip(dims, indx.values(), strict=False): # Handle chunked arrays (e.g. dask). coord = array[_dim]._variable.to_base_variable() if is_chunked_array(array.data): @@ -947,4 +952,6 @@ def _calc_idxminmax( _res.attrs = _da_idx.attrs res[_dim] = _res + if len(dims) == 1: + res = res[dims[0]] return res From b39db7d8c8c6415a4d4b4e8f44319a9b962db8e8 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Fri, 5 Sep 2025 15:44:34 -0400 Subject: [PATCH 3/7] do not cast to a list of dims --- xarray/computation/computation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/xarray/computation/computation.py b/xarray/computation/computation.py index 074254a0d3c..703cfca22dd 100644 --- a/xarray/computation/computation.py +++ b/xarray/computation/computation.py @@ -906,7 +906,9 @@ def _calc_idxminmax( # The dim is not specified and ambiguous. Don't guess. raise ValueError("Must supply 'dim' argument for multidimensional arrays") - dims = [dim] if isinstance(dim, str) else list(dim) + dim_is_str = isinstance(dim, str) + # Standardize to an iterable format + dims = [dim] if dim_is_str else dim for _dim in dims: if _dim not in array.dims: @@ -929,8 +931,8 @@ def _calc_idxminmax( # This will run argmin or argmax. indx = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna) # Force dictionary format in case of single dim so that we can iterate over it in for loop below - if len(dims) == 1: - indx = {dims[0]: indx} + if dim_is_str: + indx = {dim: indx} res = {} for _dim, _da_idx in zip(dims, indx.values(), strict=False): @@ -952,6 +954,6 @@ def _calc_idxminmax( _res.attrs = _da_idx.attrs res[_dim] = _res - if len(dims) == 1: - res = res[dims[0]] + if dim_is_str: + return res[dim] return res From a1cd3e798f223c4031b4d86b0ddf6ceaa1e1e763 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sat, 6 Sep 2025 19:17:15 -0400 Subject: [PATCH 4/7] add test_idxmin_dim for TestReduce2D --- xarray/computation/computation.py | 4 +- xarray/tests/test_dataarray.py | 62 +++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/xarray/computation/computation.py b/xarray/computation/computation.py index 703cfca22dd..850905c2f05 100644 --- a/xarray/computation/computation.py +++ b/xarray/computation/computation.py @@ -897,7 +897,9 @@ def _calc_idxminmax( if not array.ndim: raise ValueError("This function does not apply for scalars") - if dim is not None: + if dim is Ellipsis: + dim = array.dims + elif 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 diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 41604ef00ee..d7b01a6c844 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6065,6 +6065,68 @@ def test_argmax_dim( for key in expected3: assert_identical(result3[key], expected3[key]) + @pytest.mark.filterwarnings( + "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" + ) + def test_idxmin_dim( + self, + x: np.ndarray, + minindex: list[int | float], + maxindex: list[int | float], + nanindex: list[int | None], + ) -> None: + ar = xr.DataArray( + x, + dims=["y", "x"], + coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, + attrs=self.attrs, + ) + + coordarr = xr.DataArray(ar.coords["x"].data, dims=["x"]) + + if np.isnan(minindex).any(): + return + + expected0list = [coordarr.isel(x=indi, drop=True) for indi in minindex] + expected0 = {"x": xr.concat(expected0list, dim=ar.coords["y"])} + expected0["x"].name = "x" + + result0 = ar.idxmin(dim=["x"]) + for key in expected0: + assert_identical(result0[key], expected0[key]) + + result1 = ar.idxmin(dim=["x"], keep_attrs=True) + expected1 = deepcopy(expected0) + expected1["x"].attrs = self.attrs + for key in expected1: + assert_identical(result1[key], expected1[key]) + + minindex = [ + x if y is None or ar.dtype.kind == "O" else y + for x, y in zip(minindex, nanindex, strict=True) + ] + expected2list = [coordarr.isel(x=indi, drop=True) for indi in minindex] + expected2 = {"x": xr.concat(expected2list, dim=ar.coords["y"])} + expected2["x"].name = "x" + expected2["x"].attrs = {} + + result2 = ar.idxmin(dim=["x"], skipna=False) + + for key in expected2: + assert_identical(result2[key], expected2[key]) + + result3 = ar.idxmin(...) + ar_expected0 = ar.sel(expected0) + expected3 = { + "y": ar_expected0.idxmin(), + "x": ar.coords["x"][minindex[ar_expected0.argmin().item()]].reset_coords( + drop=True + ), + } + + for key in expected3: + assert_identical(result3[key], expected3[key]) + @pytest.mark.parametrize( "x, minindices_x, minindices_y, minindices_z, minindices_xy, " From af6f8e699d881ea39a627c11fa3c031a1ea01dda Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sat, 6 Sep 2025 19:19:42 -0400 Subject: [PATCH 5/7] add test_idxmax_dim for TestReduce2D --- xarray/tests/test_dataarray.py | 62 ++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index d7b01a6c844..a8dc987be4c 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6127,6 +6127,68 @@ def test_idxmin_dim( for key in expected3: assert_identical(result3[key], expected3[key]) + @pytest.mark.filterwarnings( + "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" + ) + def test_idxmax_dim( + self, + x: np.ndarray, + minindex: list[int | float], + maxindex: list[int | float], + nanindex: list[int | None], + ) -> None: + ar = xr.DataArray( + x, + dims=["y", "x"], + coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, + attrs=self.attrs, + ) + + coordarr = xr.DataArray(ar.coords["x"].data, dims=["x"]) + + if np.isnan(maxindex).any(): + return + + expected0list = [coordarr.isel(x=indi, drop=True) for indi in maxindex] + expected0 = {"x": xr.concat(expected0list, dim=ar.coords["y"])} + expected0["x"].name = "x" + + result0 = ar.idxmax(dim=["x"]) + for key in expected0: + assert_identical(result0[key], expected0[key]) + + result1 = ar.idxmax(dim=["x"], keep_attrs=True) + expected1 = deepcopy(expected0) + expected1["x"].attrs = self.attrs + for key in expected1: + assert_identical(result1[key], expected1[key]) + + maxindex = [ + x if y is None or ar.dtype.kind == "O" else y + for x, y in zip(maxindex, nanindex, strict=True) + ] + expected2list = [coordarr.isel(x=indi, drop=True) for indi in maxindex] + expected2 = {"x": xr.concat(expected2list, dim=ar.coords["y"])} + expected2["x"].name = "x" + expected2["x"].attrs = {} + + result2 = ar.idxmax(dim=["x"], skipna=False) + + for key in expected2: + assert_identical(result2[key], expected2[key]) + + result3 = ar.idxmax(...) + ar_expected0 = ar.sel(expected0) + expected3 = { + "y": ar_expected0.idxmax(), + "x": ar.coords["x"][maxindex[ar_expected0.argmax().item()]].reset_coords( + drop=True + ), + } + + for key in expected3: + assert_identical(result3[key], expected3[key]) + @pytest.mark.parametrize( "x, minindices_x, minindices_y, minindices_z, minindices_xy, " From 22d56fe05c97258c72a6ceb2d832dfb7fb3c10aa Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sat, 6 Sep 2025 20:20:39 -0400 Subject: [PATCH 6/7] add test_idxmin_dim for TestReduce3D --- xarray/tests/test_dataarray.py | 233 +++++++++++++++++++++++++++++++++ 1 file changed, 233 insertions(+) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index a8dc987be4c..a68c4eed84a 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6814,6 +6814,239 @@ def test_argmax_dim( for key in expected13: assert_identical(result13[key], expected13[key]) + def test_idxmin_dim( + self, + x: np.ndarray, + minindices_x: dict[str, np.ndarray], + minindices_y: dict[str, np.ndarray], + minindices_z: dict[str, np.ndarray], + minindices_xy: dict[str, np.ndarray], + minindices_xz: dict[str, np.ndarray], + minindices_yz: dict[str, np.ndarray], + minindices_xyz: dict[str, np.ndarray], + maxindices_x: dict[str, np.ndarray], + maxindices_y: dict[str, np.ndarray], + maxindices_z: dict[str, np.ndarray], + maxindices_xy: dict[str, np.ndarray], + maxindices_xz: dict[str, np.ndarray], + maxindices_yz: dict[str, np.ndarray], + maxindices_xyz: dict[str, np.ndarray], + nanindices_x: dict[str, np.ndarray], + nanindices_y: dict[str, np.ndarray], + nanindices_z: dict[str, np.ndarray], + nanindices_xy: dict[str, np.ndarray], + nanindices_xz: dict[str, np.ndarray], + nanindices_yz: dict[str, np.ndarray], + nanindices_xyz: dict[str, np.ndarray], + ) -> None: + ar = xr.DataArray( + x, + dims=["x", "y", "z"], + coords={ + "x": np.arange(x.shape[0]) * 4, + "y": 1 - np.arange(x.shape[1]), + "z": 2 + 3 * np.arange(x.shape[2]), + }, + attrs=self.attrs, + ) + + for inds in [ + minindices_x, + minindices_y, + minindices_z, + minindices_xy, + minindices_xz, + minindices_yz, + minindices_xyz, + ]: + if np.array([np.isnan(i) for i in inds.values()]).any(): + return + + result0 = ar.idxmin(dim=["x"]) + assert isinstance(result0, dict) + expected0 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("y", "z"), name=key) + for key, value in minindices_x.items() + } + for key in expected0: + assert_identical(result0[key].drop_vars(["y", "z"]), expected0[key]) + + result1 = ar.idxmin(dim=["y"]) + assert isinstance(result1, dict) + expected1 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x", "z"), name=key) + for key, value in minindices_y.items() + } + for key in expected1: + assert_identical(result1[key].drop_vars(["x", "z"]), expected1[key]) + + result2 = ar.idxmin(dim=["z"]) + assert isinstance(result2, dict) + expected2 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x", "y"), name=key) + for key, value in minindices_z.items() + } + for key in expected2: + assert_identical(result2[key].drop_vars(["x", "y"]), expected2[key]) + + result3 = ar.idxmin(dim=("x", "y")) + assert isinstance(result3, dict) + expected3 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("z"), name=key) + for key, value in minindices_xy.items() + } + for key in expected3: + assert_identical(result3[key].drop_vars("z"), expected3[key]) + + result4 = ar.idxmin(dim=("x", "z")) + assert isinstance(result4, dict) + expected4 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("y"), name=key) + for key, value in minindices_xz.items() + } + for key in expected4: + assert_identical(result4[key].drop_vars("y"), expected4[key]) + + result5 = ar.idxmin(dim=("y", "z")) + assert isinstance(result5, dict) + expected5 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x"), name=key) + for key, value in minindices_yz.items() + } + for key in expected5: + assert_identical(result5[key].drop_vars("x"), expected5[key]) + + result6 = ar.idxmin(...) + assert isinstance(result6, dict) + expected6 = { + key: xr.DataArray(ar.coords[key].values[value], name=key) + for key, value in minindices_xyz.items() + } + for key in expected6: + assert_identical(result6[key], expected6[key]) + """ + + minindices_x = { + key: xr.where( + nanindices_x[key] == None, # noqa: E711 + minindices_x[key], + nanindices_x[key], + ) + for key in minindices_x + } + expected7 = { + key: xr.DataArray(value, dims=("y", "z")) + for key, value in minindices_x.items() + } + + result7 = ar.idxmin(dim=["x"], skipna=False) + assert isinstance(result7, dict) + for key in expected7: + assert_identical(result7[key].drop_vars(["y", "z"]), expected7[key]) + + minindices_y = { + key: xr.where( + nanindices_y[key] == None, # noqa: E711 + minindices_y[key], + nanindices_y[key], + ) + for key in minindices_y + } + expected8 = { + key: xr.DataArray(value, dims=("x", "z")) + for key, value in minindices_y.items() + } + + result8 = ar.idxmin(dim=["y"], skipna=False) + assert isinstance(result8, dict) + for key in expected8: + assert_identical(result8[key].drop_vars(["x", "z"]), expected8[key]) + + minindices_z = { + key: xr.where( + nanindices_z[key] == None, # noqa: E711 + minindices_z[key], + nanindices_z[key], + ) + for key in minindices_z + } + expected9 = { + key: xr.DataArray(value, dims=("x", "y")) + for key, value in minindices_z.items() + } + + result9 = ar.idxmin(dim=["z"], skipna=False) + assert isinstance(result9, dict) + for key in expected9: + assert_identical(result9[key].drop_vars(["x", "y"]), expected9[key]) + + minindices_xy = { + key: xr.where( + nanindices_xy[key] == None, # noqa: E711 + minindices_xy[key], + nanindices_xy[key], + ) + for key in minindices_xy + } + expected10 = { + key: xr.DataArray(value, dims="z") for key, value in minindices_xy.items() + } + + result10 = ar.idxmin(dim=("x", "y"), skipna=False) + assert isinstance(result10, dict) + for key in expected10: + assert_identical(result10[key].drop_vars("z"), expected10[key]) + + minindices_xz = { + key: xr.where( + nanindices_xz[key] == None, # noqa: E711 + minindices_xz[key], + nanindices_xz[key], + ) + for key in minindices_xz + } + expected11 = { + key: xr.DataArray(value, dims="y") for key, value in minindices_xz.items() + } + + result11 = ar.idxmin(dim=("x", "z"), skipna=False) + assert isinstance(result11, dict) + for key in expected11: + assert_identical(result11[key].drop_vars("y"), expected11[key]) + + minindices_yz = { + key: xr.where( + nanindices_yz[key] == None, # noqa: E711 + minindices_yz[key], + nanindices_yz[key], + ) + for key in minindices_yz + } + expected12 = { + key: xr.DataArray(value, dims="x") for key, value in minindices_yz.items() + } + + result12 = ar.idxmin(dim=("y", "z"), skipna=False) + assert isinstance(result12, dict) + for key in expected12: + assert_identical(result12[key].drop_vars("x"), expected12[key]) + + minindices_xyz = { + key: xr.where( + nanindices_xyz[key] == None, # noqa: E711 + minindices_xyz[key], + nanindices_xyz[key], + ) + for key in minindices_xyz + } + expected13 = {key: xr.DataArray(value) for key, value in minindices_xyz.items()} + + result13 = ar.idxmin(..., skipna=False) + assert isinstance(result13, dict) + for key in expected13: + assert_identical(result13[key], expected13[key]) + """ + class TestReduceND(TestReduce): @pytest.mark.parametrize("op", ["idxmin", "idxmax"]) From f3cec881b1662dd6fb01e5e92d548a41bdb68528 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sat, 6 Sep 2025 20:39:52 -0400 Subject: [PATCH 7/7] add test_idxmax_dim for TestReduce3D --- xarray/tests/test_dataarray.py | 233 +++++++++++++++++++++++++++++++++ 1 file changed, 233 insertions(+) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index a68c4eed84a..baebfdd264f 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -7047,6 +7047,239 @@ def test_idxmin_dim( assert_identical(result13[key], expected13[key]) """ + def test_idxmax_dim( + self, + x: np.ndarray, + minindices_x: dict[str, np.ndarray], + minindices_y: dict[str, np.ndarray], + minindices_z: dict[str, np.ndarray], + minindices_xy: dict[str, np.ndarray], + minindices_xz: dict[str, np.ndarray], + minindices_yz: dict[str, np.ndarray], + minindices_xyz: dict[str, np.ndarray], + maxindices_x: dict[str, np.ndarray], + maxindices_y: dict[str, np.ndarray], + maxindices_z: dict[str, np.ndarray], + maxindices_xy: dict[str, np.ndarray], + maxindices_xz: dict[str, np.ndarray], + maxindices_yz: dict[str, np.ndarray], + maxindices_xyz: dict[str, np.ndarray], + nanindices_x: dict[str, np.ndarray], + nanindices_y: dict[str, np.ndarray], + nanindices_z: dict[str, np.ndarray], + nanindices_xy: dict[str, np.ndarray], + nanindices_xz: dict[str, np.ndarray], + nanindices_yz: dict[str, np.ndarray], + nanindices_xyz: dict[str, np.ndarray], + ) -> None: + ar = xr.DataArray( + x, + dims=["x", "y", "z"], + coords={ + "x": np.arange(x.shape[0]) * 4, + "y": 1 - np.arange(x.shape[1]), + "z": 2 + 3 * np.arange(x.shape[2]), + }, + attrs=self.attrs, + ) + + for inds in [ + maxindices_x, + maxindices_y, + maxindices_z, + maxindices_xy, + maxindices_xz, + maxindices_yz, + maxindices_xyz, + ]: + if np.array([np.isnan(i) for i in inds.values()]).any(): + return + + result0 = ar.idxmax(dim=["x"]) + assert isinstance(result0, dict) + expected0 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("y", "z"), name=key) + for key, value in maxindices_x.items() + } + for key in expected0: + assert_identical(result0[key].drop_vars(["y", "z"]), expected0[key]) + + result1 = ar.idxmax(dim=["y"]) + assert isinstance(result1, dict) + expected1 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x", "z"), name=key) + for key, value in maxindices_y.items() + } + for key in expected1: + assert_identical(result1[key].drop_vars(["x", "z"]), expected1[key]) + + result2 = ar.idxmax(dim=["z"]) + assert isinstance(result2, dict) + expected2 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x", "y"), name=key) + for key, value in maxindices_z.items() + } + for key in expected2: + assert_identical(result2[key].drop_vars(["x", "y"]), expected2[key]) + + result3 = ar.idxmax(dim=("x", "y")) + assert isinstance(result3, dict) + expected3 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("z"), name=key) + for key, value in maxindices_xy.items() + } + for key in expected3: + assert_identical(result3[key].drop_vars("z"), expected3[key]) + + result4 = ar.idxmax(dim=("x", "z")) + assert isinstance(result4, dict) + expected4 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("y"), name=key) + for key, value in maxindices_xz.items() + } + for key in expected4: + assert_identical(result4[key].drop_vars("y"), expected4[key]) + + result5 = ar.idxmax(dim=("y", "z")) + assert isinstance(result5, dict) + expected5 = { + key: xr.DataArray(ar.coords[key].values[value], dims=("x"), name=key) + for key, value in maxindices_yz.items() + } + for key in expected5: + assert_identical(result5[key].drop_vars("x"), expected5[key]) + + result6 = ar.idxmax(...) + assert isinstance(result6, dict) + expected6 = { + key: xr.DataArray(ar.coords[key].values[value], name=key) + for key, value in maxindices_xyz.items() + } + for key in expected6: + assert_identical(result6[key], expected6[key]) + """ + + maxindices_x = { + key: xr.where( + nanindices_x[key] == None, # noqa: E711 + maxindices_x[key], + nanindices_x[key], + ) + for key in maxindices_x + } + expected7 = { + key: xr.DataArray(value, dims=("y", "z")) + for key, value in maxindices_x.items() + } + + result7 = ar.idxmax(dim=["x"], skipna=False) + assert isinstance(result7, dict) + for key in expected7: + assert_identical(result7[key].drop_vars(["y", "z"]), expected7[key]) + + maxindices_y = { + key: xr.where( + nanindices_y[key] == None, # noqa: E711 + maxindices_y[key], + nanindices_y[key], + ) + for key in maxindices_y + } + expected8 = { + key: xr.DataArray(value, dims=("x", "z")) + for key, value in maxindices_y.items() + } + + result8 = ar.idxmax(dim=["y"], skipna=False) + assert isinstance(result8, dict) + for key in expected8: + assert_identical(result8[key].drop_vars(["x", "z"]), expected8[key]) + + maxindices_z = { + key: xr.where( + nanindices_z[key] == None, # noqa: E711 + maxindices_z[key], + nanindices_z[key], + ) + for key in maxindices_z + } + expected9 = { + key: xr.DataArray(value, dims=("x", "y")) + for key, value in maxindices_z.items() + } + + result9 = ar.idxmax(dim=["z"], skipna=False) + assert isinstance(result9, dict) + for key in expected9: + assert_identical(result9[key].drop_vars(["x", "y"]), expected9[key]) + + maxindices_xy = { + key: xr.where( + nanindices_xy[key] == None, # noqa: E711 + maxindices_xy[key], + nanindices_xy[key], + ) + for key in maxindices_xy + } + expected10 = { + key: xr.DataArray(value, dims="z") for key, value in maxindices_xy.items() + } + + result10 = ar.idxmax(dim=("x", "y"), skipna=False) + assert isinstance(result10, dict) + for key in expected10: + assert_identical(result10[key].drop_vars("z"), expected10[key]) + + maxindices_xz = { + key: xr.where( + nanindices_xz[key] == None, # noqa: E711 + maxindices_xz[key], + nanindices_xz[key], + ) + for key in maxindices_xz + } + expected11 = { + key: xr.DataArray(value, dims="y") for key, value in maxindices_xz.items() + } + + result11 = ar.idxmax(dim=("x", "z"), skipna=False) + assert isinstance(result11, dict) + for key in expected11: + assert_identical(result11[key].drop_vars("y"), expected11[key]) + + maxindices_yz = { + key: xr.where( + nanindices_yz[key] == None, # noqa: E711 + maxindices_yz[key], + nanindices_yz[key], + ) + for key in maxindices_yz + } + expected12 = { + key: xr.DataArray(value, dims="x") for key, value in maxindices_yz.items() + } + + result12 = ar.idxmax(dim=("y", "z"), skipna=False) + assert isinstance(result12, dict) + for key in expected12: + assert_identical(result12[key].drop_vars("x"), expected12[key]) + + maxindices_xyz = { + key: xr.where( + nanindices_xyz[key] == None, # noqa: E711 + maxindices_xyz[key], + nanindices_xyz[key], + ) + for key in maxindices_xyz + } + expected13 = {key: xr.DataArray(value) for key, value in maxindices_xyz.items()} + + result13 = ar.idxmax(..., skipna=False) + assert isinstance(result13, dict) + for key in expected13: + assert_identical(result13[key], expected13[key]) + """ + class TestReduceND(TestReduce): @pytest.mark.parametrize("op", ["idxmin", "idxmax"])