在嵌套循环中使用 dask 进行循环并行化

Question

我只是在学习使用 dask，并在这个论坛上阅读了许多与 Dask 和 for 循环相关的主题。但我仍然不清楚如何将这些解决方案应用于我的问题。我正在处理与（时间、深度、位置）有关的气候数据。 “位置”坐标是一个线性索引，因此每个值对应一个唯一的（经度、纬度）。假设 var1 和 var2 是两个输入变量，我在下面展示了我正在尝试做的事情的基本框架。我想对位置参数“nxy”进行并行化，因为我的计算可以在不同的位置同时进行。

for loc in range(0,nxy):           # nxy = total no. of locations
    for it in range(0,ntimes):    

       out1 = expression1 involving ( var1(loc), var2(it,loc) )
       out2 = expression2 involving ( var1(loc), var2(it,loc) )
       # <a dozen more output variables>

我的问题：

(i) 许多说明“延迟”使用的示例显示了“延迟（函数）（arg）”之类的内容。就我而言，我没有太多（如果有的话）函数，但有很多表达式。如果'delayed'只在函数级别操作，是否应该把每个表达式都转换成一个函数，并在前面加上'delayed'？

(ii) 我是否应该将上面显示的整个 for 循环包装在一个函数中，然后使用“延迟”调用该函数？我尝试做这样的事情，但可能做不正确，因为与不使用 dask 相比，我没有得到任何加速。这是我所做的：

def test_dask(n):
 for loc in range(0,n):
   # same code as before
 return var1  # just returning one variable for now

var1=delayed(tast_dask)(nxy)
var1.compute()

感谢您的帮助。

Answer 1

每个延迟的任务都会增加大约 1 毫秒的开销。因此，如果您的表达速度很慢（也许您正在调用其他一些昂贵的函数），那么是的 dask.delayed 可能是一个不错的选择。如果没有，那么您可能应该寻找其他地方。

特别是，看起来您只是在遍历几个数组并逐个元素地操作。请注意，Python 在这方面非常慢。您可能根本不想使用 Dask，而是尝试以下方法之一：

1. 找到一些巧妙的方法来使用 Numpy 表达式重写您的计算
2. 使用 Numba

另外，考虑到您使用的术语，如 lat/lon/depth，Xarray 可能对您来说是一个很好的项目。