如何解决xarray:处理来自ufunc的多个不同维度的返回值
我有一个对一维数组(时间序列)起作用的函数,并返回该数组的修改版本以及一个系数。
我正在使用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>
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。