xarray：处理来自 ufunc 的不同维度的多个 return 值

Question

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

我正在使用 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、...等等，但我无法得到我想要的结果。

Answer 1

感谢@Maximilian 的有用评论，我有了一个使用包装器的可行解决方案。

由于直接从 apply_ufunc 中获得组合数据集似乎不太可能，我将 post 这作为答案（我会为某人 post 保留可能性改进）。

正在使用

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>

xarray：处理来自 ufunc 的不同维度的多个 return 值

xarray: dealing with multiple return values of different dimensions from ufunc

python

numpy-ufunc

dask

python-xarray