У меня 4D xarray Dataset. Я хочу провести линейную регрессию между двумя переменными на конкретном измерении (здесь время) и сохранить параметры регрессии в трехмерном массиве (остальные измерения). Мне удалось получить результаты, которые я хочу, используя этот серийный код, но он довольно медленный:
# 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: переместить предыдущее редактирование в ответ
Можно применить scipy.stats.linregress
(и другие не-ufunc) к набору данных apply_ufunc()
используя 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'