xarray：处理来自ufunc的不同维度的多个返回值答案

【问题标题】：xarray: dealing with multiple return values of different dimensions from ufuncxarray：处理来自ufunc的不同维度的多个返回值
【发布时间】：2021-02-23 17:06:27
【问题描述】：

我有一个函数适用于一维数组（时间序列）并返回数组的修改版本，加上一个系数。

我正在使用 xarray.apply_ufunc 将此函数并行应用于我的 3D dask 数组。

所需的输出是一个 xarray 数据集，其中包含修改后的 3D 数组，以及一个包含每个时间序列系数的新变量（因此本质上是一个 2D 数组）。

在下面的示例中，函数 fun 通过返回输入数据乘以 2 加上一个随机整数来模拟真实函数的行为。

import xarray as xr
import numpy as np

# function
def fun(x):

    return x*2, np.random.randint(1000, size=(1,))

# test dataset
ds = xr.tutorial.open_dataset("air_temperature", chunks={"time": -1}).isel(time=slice(100,200))

print(ds)

# <xarray.Dataset>
# Dimensions:  (lat: 25, lon: 53, time: 100)
# Coordinates:
#   * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
#   * lon      (lon) float32 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0
#   * time     (time) datetime64[ns] 1800-01-01 1800-01-01 ... 1800-01-01
# Data variables:
#     air      (time, lat, lon) float32 dask.array<chunksize=(100, 25, 53), meta=np.ndarray>
# Attributes:
#     Conventions:  COARDS
#     title:        4x daily NMC reanalysis (1948)
#     description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
#     platform:     Model
#     references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...

现在将其应用于数据集：


np.random.seed(42)

ds2 = xr.apply_ufunc(
    fun,
    ds,
    input_core_dims=[['time']],
    output_core_dims=[["time"],[]],
    vectorize=True,
    dask="parallelized"
)

这几乎产生了我需要的东西。我得到一个包含两个数据集的元组，一个是 3D 数组，另一个是带有系数的 2D 数组。

(<xarray.Dataset>
 Dimensions:  (lat: 25, lon: 53, time: 100)
 Coordinates:
   * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
   * lon      (lon) float32 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0
   * time     (time) datetime64[ns] 1800-01-01 1800-01-01 ... 1800-01-01
 Data variables:
     air      (lat, lon, time) float32 496.78 496.78 493.59998 ... 0.0 0.0 0.0,
 <xarray.Dataset>
 Dimensions:  (lat: 25, lon: 53)
 Coordinates:
   * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
   * lon      (lon) float32 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0
 Data variables:
     air      (lat, lon) int64 120 721 93 288 135 478 ... 380 497 881 102 485 814)

但我真正想要的是一个组合数据集，之后可以使用dask 将其作为一个整体进行计算。我试过修改input_core_dims、output_core_dims、output_sizes，...等，但我无法得到我想要的结果。

【问题讨论】：

感谢您提供的清晰示例！就预期输出而言——如何组合数据集？他们目前有一个重叠的air 变量
@Maximilian 感谢您的评论！第二个数据集 (2d) 中的 air 变量可以重命名，这就是它从 apply_ufunc 中出来的方式。它实际上是源自air 的系数。理想情况下，组合发生在apply_ufunc，返回值是具有 2 个变量的单个数据集。
我明白了！我认为最好的方法是进行这种合并的包装器函数——我认为没有办法在apply_ufunc内做到这一点，并且包装器相当普遍：xarray.pydata.org/en/stable/generated/xarray.apply_ufunc.html跨度>

标签： python dask python-xarray numpy-ufunc

【解决方案1】：

感谢@Maximilian 的帮助，我有一个使用包装器的工作解决方案。

由于直接从 apply_ufunc 获得组合数据集似乎不太可能，因此我将发布此作为答案（我会为发布改进的人保持开放的可能性）。

使用

def wrap(ds):

    ds2, coeff = xr.apply_ufunc(
        fun,
        ds.air,
        input_core_dims=[['time']],
        output_core_dims=[["time"],[]],
        vectorize=True,
        dask="parallelized"
    )
    
    ds_out = ds2.to_dataset()
    ds_out["coeff"] = coeff

    return ds_out

给我

wrap(ds)

#<xarray.Dataset>
#Dimensions:  (lat: 25, lon: 53, time: 100)
#Coordinates:
#  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
#  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
#  * time     (time) datetime64[ns] 2013-01-26 ... 2013-02-19T18:00:00
#Data variables:
#    air      (lat, lon, time) float32 dask.array<chunksize=(25, 53, 100), #meta=np.ndarray>
#    coeff    (lat, lon) int64 dask.array<chunksize=(25, 53), meta=np.ndarray>

【讨论】：