Question

我有一个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：将先前的编辑移至答案

Answer 1

LCT的上一个答案涵盖了这里应该说的大部分内容，但是我认为可以将dask='parallelized'与来自{{ 1}}。

这里的技巧是将多个输出堆叠到一个数组中，然后输出，您还必须使用scipy.stats.linregress kwarg来指定output_core_dims调用中的DataArray输出将具有现在增加了尺寸：

apply_ufunc()

def new_linregress(x, y):
    # Wrapper around scipy linregress to use in apply_ufunc
    slope, intercept, r_value, p_value, std_err = stats.linregress(x, y)
    return np.array([slope, intercept, r_value, p_value, std_err])

注意事项：如果您拥有# return a new DataArray stats = xr.apply_ufunc(new_linregress, ds[x], ds[y], input_core_dims=[['year'], ['year']], output_core_dims=[["parameter"]], vectorize=True, dask="parallelized", output_dtypes=['float64'], output_sizes={"parameter": 5}, )，则此方法当前仅适用于dask='parallelized'，但如果您有类似{{1} }。请查看此Github issue以获得更多详细信息。

希望有帮助！

编辑：我们已通知您，dask<2.0的问题dask='allowed'已得到纠正！因此，现在可以使用dask<2.0加快速度。：）

Answer 2

通过传递scipy.stats.linregress可以使用apply_ufunc()将vectorize=True（和其他非函数）应用于xarray数据集，如下所示：

# 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倍。

但是，仍然无法像scipy.stats.linregress那样通过多个输出实现dask的并行化：

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

使用xarray，如何并行处理多维数据集上的一维操作？

2 个答案: