时间戳

首页 > 解决方案 > 使用 Xarray 在 DataSet 中迭代添加 2D DataArrays 而不会丢失维度的 Pythonic 方法？

python - 使用 Xarray 在 DataSet 中迭代添加 2D DataArrays 而不会丢失维度的 Pythonic 方法？

问题描述

我有一堆由可变数量的 2D DataArrays 组成的 xarray 数据集，大小都相同（纬度/经度暗淡）。我想得到数据集中一些DataArrays 的总和，而不会丢失任何维度.

我找到了一种方法来做到这一点，但似乎应该有一个内置的函数/更 Pythonic 的方式来做到这一点。可以在单行 for 循环或使用迭代中完成吗？似乎大多数求和函数都会迫使你失去一个维度。

# Build a fake dataset that is way, way smaller than what I have, with lat, lon as dims: 
ds=xr.Dataset()
lat=np.arange(20, 50, 2)
lon=np.arange(-130, -60, 2.5)

ds['emis_NO']= xr.DataArray(np.random.random([len(lat), len(lon)]), dims=['lat','lon'])
ds['emis_NO2']= xr.DataArray(np.random.random([len(lat), len(lon)]), dims=['lat','lon'])
ds['emis_ISOP']= xr.DataArray(np.random.random([len(lat), len(lon)]), dims=['lat','lon'])
ds['emis_ACET']= xr.DataArray(np.random.random([len(lat), len(lon)]), dims=['lat','lon'])

# Get a list of the DataArrays I want to Sum in this DataSet, ds, matching a 
# string naming convention (Sum emissions of NOx = NO + NO2). 
vars2sum=  [x for x in ds.data_vars if ('emis_NO' in x)]

# What follows is my solution that works, that I'd like to condense 
# if possible in a more pythonic way: 
total= ds[vars2sum[0]].values * 0 # create something right size, fill w/ 0s 

for darr in vars2sum:             # Loop over variables to sum. 
   total= total+ ds[darr].values  # iteratively update total.

所需的输出相当于手动添加 DataArrays，如下所示（假设您都知道它们）。

total= ds.emis_NO.values + ds.emis_NO2.values

我在一个函数中进行迭代，因为我在给定 DataSet 中求和的 DataArrays 的数量（通过我的 str 匹配捕获）因数据集而异。

标签： pythonnumpymultidimensional-arraypython-xarrayiterable

解决方案

to_array您可以首先使用（默认情况下，新维度将命名为）沿新维度连接所选变量'variable'。然后，沿该维度求和：

ds[vars2sum].to_array().sum("variable")

或者，您可以使用列表推导：

sum([ds[var] for var in vars2sum])

推荐阅读