首页 文章

使用xarray,如何在多维数据集上并行化1D操作?

提问于
浏览
1

我有一个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并行化),但我没有设法使参数正确 .

以下问题是相关的,但没有答案:

Edit 2: move previous edit to an answer

1 回答

  • 0

    可以使用 apply_ufunc()scipy.stats.linregress (和其他非ufuncs)应用于xarray数据集,方法是传递 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'

相关问题