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 for循环很慢。通常，加速这样的代码的第一步是...

找到一种巧妙的方式将所有这些内容编写为矢量化的numpy表达式，而无需Python for循环
使用Numba

使用Dask（或其他有效方式）并行化循环

1 个答案: