xarray 在具有不同变量名称的两个数据集之间应用函数

如何解决xarray 在具有不同变量名称的两个数据集之间应用函数

我有两个类似的 xarray 数据集（来自 grib 文件）A 和 B，其结构如下：

<xarray.Dataset>
Dimensions:            (latitude: 198,longitude: 318,step: 42)
Coordinates:
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 00:00:00 ... 1 days 17:00:00
    heightAboveGround  int64 ...
  * latitude           (latitude) float64 48.3 48.27 48.25 ... 43.43 43.4 43.38
  * longitude          (longitude) float64 -7.2 -7.175 -7.15 ... 0.675 0.7 0.725
    valid_time         (step) datetime64[ns] ...
Data variables:
    u10                (step,latitude,longitude) float32 ...
Attributes:
    GRIB_edition:            2
    GRIB_centre:             lfpw
    GRIB_centreDescription:  french Weather Service - Toulouse
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             french Weather Service - Toulouse
    history:                 2021-03-02T18:34:18 GRIB to CDM+CF via cfgrib-0....

数据集 A 有一个名为 u10 的变量，B 有一个名为 v10 的变量。它们是我要计算的风向的两个组成部分。这些值是为每个空间点（纬度和经度）和时间（42 步）设置的。

为了记录，这里描述了 Python 中相应的向量化计算公式：How to calculate wind direction from U and V wind components in R

dir = np.mod(180+np.rad2deg(np.arctan2(U,V)),360)

但是我很难将这个函数应用到我的数据集上。我尝试使用以下模式（在此处找到 https://github.com/pydata/xarray/issues/900）但没有成功：

ds1 = xr.open_dataset("u_component.grib",engine='cfgrib')
ds2 = xr.open_dataset("v_component.grib",engine='cfgrib')
combined = xr.Dataset({'u_component': ds1,'v_composante': ds2})
combined.groupby('step').apply(compute_direction)

首先，我不知道如何将这两个变量组合在一个数据集中（但也许没有必要）。由于多个维度，我对 compute_direction 函数定义感到非常困惑。

有人有一些关于如何进行的示例或提示吗？