使用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循环包装在一个函数中，然后使用“delayed”调用该函数？我尝试过这样做，但可能做得不对，因为与不使用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

每个延迟的任务都会增加大约1ms的开销。因此，如果您的表达式很慢（可能您正在调用其他昂贵的函数），那么yes dask.delayed可能是一个很好的选择。如果没有，那么你应该去别处看看

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

1. 找到一些聪明的方法用Numpy表达式重写计算
2. 用麻木

此外，考虑到您使用的术语（如lat/lon/depth），Xarray可能是一个很好的项目