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

编程问答 2022-06-03

问题描述

我有一个对一维数组（时间序列）起作用的函数，并返回该数组的修改版本以及一个系数。

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

解决方法

由于@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>
dasknumpy-ufuncpythonpython-xarray