有没有更快的方法来汇总 Xarray 数据集变量？答案

【问题标题】：Is there a faster way to sum Xarray dataset variables?有没有更快的方法来汇总 Xarray 数据集变量？
【发布时间】：2021-03-16 20:21:58
【问题描述】：

这是我的第一个堆栈交换问题，所以我希望我做对了。

我正在尝试将数据集中的几个 xarray 变量相加。每个变量具有相同的维度。代码基本上是这样的：

def add_variables(xarray_dataset, listofvars):
    data = 0
    for var in listofvars:
        data = data + dset[var][:,-1,:] # slice of each variable
    return data 

summed_variables = add_variables(dset, ['varname1, varname2'])

但是，这需要很长时间才能运行。有没有人建议一种更快的方法来解决这个问题？谢谢！

【问题讨论】：

你试过xarray_dataset.sum()吗？如果这不正确，您能否发布一个完全可重现的示例？

标签： python dataset python-xarray

【解决方案1】：

您可以使用to_array 方法沿新维度（默认命名为“变量”）堆叠变量，然后在该维度上求和。如有必要，您可以预先选择变量并对其进行切片。

import numpy as np
import xarray as xr

# Create dummy dataset
ds = xr.Dataset(
    {var: (("x", "y", "z"), np.random.rand(5, 3, 2)) for var in "abcde"}
)

# Sum over (a slice of some of the) variables
vars_to_sum = ["a", "c", "d"]
summed_variables = ds[vars_to_sum].isel(y=-1).to_array().sum("variable")

我认为这比您的自定义函数要容易得多，尽管在我的比较中它并不快¹：

%timeit add_variables(ds, vars_to_sum)
464 µs ± 591 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit ds[vars_to_sum].isel(y=-1).to_array().sum("variable")
660 µs ± 12.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

但是，对于这个小数据集，它们都非常快，因此差异并不明显。我不知道你的数据集是什么样子的——如果你能分享一些关于数据的更多信息以诊断性能问题，这可能会有所帮助。

¹ 请注意，我必须稍微更改您的函数才能使其运行——函数头和主体中的数据集名称不一致：

def add_variables(xarray_dataset, listofvars):
    data = 0
    for var in listofvars:
        # changed dset to xarray_dataset in the following line
        data = data + xarray_dataset[var][:,-1,:]
    return data

【讨论】：