Распараллеливание для циклов с использованием Dask (или другим эффективным способом)

Question

У меня есть функция, которая принимает набор данных xarray (аналогично pandas multi-index) и использует 4 для циклов, встроенных друг в друга, для вычисления новой переменной массива данных.

Интересно, есть ли Я могу использовать Dask, чтобы ускорить этот процесс, я новичок в этом, поэтому я не уверен.

Функция выглядит следующим образом:

def A_calc(data, thresh):

    A = np.zeros((len(data.time), len(data.lat), len(data.lon)))
    foo = xr.DataArray(A, coords=[data.time, data.lat, data.lon], 
                   dims=['time','lat', 'lon'])

    for t in tqdm(range(len(data.time))):
        for i in range(len(data.lat)):
            for j in range(2,len(data.lon)):
                for k in range(len(data.lev)):
                    if np.isnan(
                            data[dict(time=[t], lat=[i], lon=[j], lev=[k])].sigma_0.values):
                            foo[dict(time=[t], lat=[i], lon=[j])] = np.nan
                            break
                    elif abs(
                         data[dict(time=[t], lat=[i], lon=[j], lev=[k])].sigma_0.values
                         - data[dict(time=[t], lat=[i], lon=[j], lev=[1])].sigma_0.values) >= thresh:
                        foo[dict(time=[t], lat=[i], lon=[j])] = data.lev[k].values
                        break

    return foo        

Есть предложения?

Answer 1

Как сказано в комментариях, Python для циклов медленные. Как правило, первым шагом к ускорению кода, подобного этому, является либо ...

1. Найти какой-нибудь умный способ написать все это как векторизованное numpy выражение, без Python для циклов
2. Используйте Numba