时间戳

首页 > 解决方案 > 使用 xarray，如何在多维数据集上并行化一维操作？

python - 使用 xarray，如何在多维数据集上并行化一维操作？

问题描述

我有一个 4D xarray 数据集。我想在特定维度（此处为时间）上的两个变量之间进行线性回归，并将回归参数保留在 3D 数组（其余维度）中。我设法通过使用此串行代码获得了我想要的结果，但它相当慢：

# add empty arrays to store results of the regression
res_shape = tuple(v for k,v in ds[x].sizes.items() if k != 'year')
res_dims = tuple(k for k,v in ds[x].sizes.items() if k != 'year')
ds[sl] = (res_dims, np.empty(res_shape, dtype='float32'))
ds[inter] = (res_dims, np.empty(res_shape, dtype='float32'))
# Iterate in kept dimensions
for lat in ds.coords['latitude']:
    for lon in ds.coords['longitude']:
        for duration in ds.coords['duration']:
            locator = {'longitude':lon, 'latitude':lat, 'duration':duration}
            sel = ds.loc[locator]
            res = scipy.stats.linregress(sel[x], sel[y])
            ds[sl].loc[locator] = res.slope
            ds[inter].loc[locator] = res.intercept

我怎样才能加速和并行化这个操作？

我知道这apply_ufunc可能是一个选项（并且可以与 dask 并行化），但我没有设法使参数正确。

以下问题是相关的，但没有答案：

编辑 2：将先前的编辑移至答案

标签： pythondaskpython-xarray

解决方案

可以通过如下方式将scipy.stats.linregress（和其他非 ufunc）应用到 xarray 数据集：apply_ufunc()vectorize=True

# return a tuple of DataArrays
res = xr.apply_ufunc(scipy.stats.linregress, ds[x], ds[y],
        input_core_dims=[['year'], ['year']],
        output_core_dims=[[], [], [], [], []],
        vectorize=True)
# add the data to the existing dataset
for arr_name, arr in zip(array_names, res):
    ds[arr_name] = arr

虽然仍然apply_ufunc是串行的，但在这种特定情况下比循环实现快大约 36 倍。

然而，与 dask 的并行化仍然没有通过多个输出来实现，例如来自scipy.stats.linregress：

NotImplementedError：来自 apply_ufunc 的多个输出尚不支持 dask='parallelized'

推荐阅读