利用一个数据组的argmax得到另一组对应的数据

Question

描述

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, 1, 3, 4]))},
    coords={"x": [10, 10, 20, 20]},
)


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

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

$ print(max_id)

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

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

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

但是，它将选择应用于 test2

的每个 x 组

$ 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

Answer 1

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

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，所以不知道是否存在其他更有效的方法

利用一个数据组的argmax得到另一组对应的数据

Use the argmax of one data group to get the corresponding data of another group

python

pandas

python-xarray

描述

例子

预期结果