使用一个数据组的argmax得到另一组对应的数据 说明示例预期结果

问题描述

说明

ds 数据集有两个 DataArray：test1 和 test2。

test1 由 x 分组，最大值由 argmax 获得。

我想通过索引获取test2对应的位置数据。

但是，我不知道正确的方法。

示例

import xarray as xr
import numpy as np

ds = xr.Dataset(
    {"test1": (("x"),np.array([0,1,3,2])),"test2": (("x"),np.array([2,4]))},coords={"x": [10,10,20,20]},)


max_id = ds['test1'].groupby('x').apply(lambda da: da.argmax(dim='x'))

获取每个test1组中x DataArray的最大值索引的方法效果很好：

$ print(max_id)

<xarray.DataArray 'test1' (x: 2)>
array([1,0])
Coordinates:
  * x        (x) int64 10 20

我尝试使用 max_id 对按 test2 分组的 x 进行子集化：

test2_max = ds['test2'].groupby('x').apply(lambda da: da.isel(x=max_id.values))

但是，它将选择应用于每个 x 组的 test2

$ print(test2_max)

<xarray.DataArray 'test2' (x: 4)>
array([1,2,4,3])
Coordinates:
  * x        (x) int64 10 10 20 20

预期结果

$ print(test2_max)

<xarray.DataArray 'test2' (x: 2)>
array([1,3])
Coordinates:
  * x        (x) int64 10 20

解决方法

我正在玩这个，我能够做到以下几点：

import xarray as xr
import numpy as np

ds = xr.Dataset(
    {"test1": (("x"),np.array([0,1,3,2])),"test2": (("x"),np.array([1,2,4,3]))},coords={"x": [10,10,20,20]},)

test2_max = ds.groupby('x').apply(lambda ds: ds['test2'].isel(x=ds['test1'].argmax(dim='x')))

输出：

xarray.DataArray'test2'x: 2
array([2,4])
Coordinates:
x
(x)
int64
10 20
Attributes: (0)

我从未使用过 xarray，所以不知道是否存在其他更有效的方法