我有一个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 回答
可以使用
apply_ufunc()
将scipy.stats.linregress
(和其他非ufuncs)应用于xarray数据集,方法是传递vectorize=True
,如下所示:虽然仍然是串行的,但在这种特定情况下,
apply_ufunc
比循环实现快36倍 .但是,与dask的并行化仍未实现多个输出,如
scipy.stats.linregress
中的那个: