沿dask数组的轴应用函数

问题内容：

我正在分析气候模型模拟中的海洋温度数据，其中4D数据阵列（时间，深度，纬度，经度；dask_array如下所示）的形状通常为（6000、31、189、192），大小约为25GB（因此，我渴望使用dask；尝试使用numpy处理这些数组时，我一直遇到内存错误。

我需要沿着时间轴在每个级别/纬度/经度点处拟合三次多项式，并存储所得的4个系数。因此，我已经设置chunksize=(6000, 1, 1, 1)好，因此每个网格点都有一个单独的块。

这是我获取三次多项式系数的函数（time_axis轴值是在其他位置定义的全局一维numpy数组）：

def my_polyfit(data):    
    return numpy.polyfit(data.squeeze(), time_axis, 3)

（因此，在这种情况下，numpy.polyfit返回长度为4的列表）

这是我认为我需要将其应用于每个块的命令：

dask_array.map_blocks(my_polyfit, chunks=(4, 1, 1, 1), drop_axis=0, new_axis=0).compute()

现在，时间轴消失了（因此drop_axis=0），并且在它的位置（长度为4）处有了一个新的系数轴。

当我运行此命令时，我得到了IndexError: tuple index out of range，所以我想知道我在哪里/如何误解了map_blocks？

问题答案：

我怀疑如果您的函数返回与它消耗的维相同的数组，您的体验会更流畅。例如，您可以考虑如下定义函数：

def my_polyfit(data):
    return np.polyfit(data.squeeze(), ...)[:, None, None, None]

然后，您可能会忽略new_axis，drop_axis位。

在性能方面，您可能还想考虑使用更大的块大小。如果每个块有6000个数字，那么您将有超过一百万个块，这意味着与实际计算相比，您可能会花费更多的时间进行调度。通常，我会拍摄几兆字节的块。当然，增加块大小会导致映射函数变得更加复杂。

例

In [1]: import dask.array as da

In [2]: import numpy as np

In [3]: def f(b):
    return np.polyfit(b.squeeze(), np.arange(5), 3)[:, None, None, None]
   ...:

In [4]: x = da.random.random((5, 3, 3, 3), chunks=(5, 1, 1, 1))

In [5]: x.map_blocks(f, chunks=(4, 1, 1, 1)).compute()
Out[5]: 
array([[[[ -1.29058580e+02,   2.21410738e+02,   1.00721521e+01],
         [ -2.22469851e+02,  -9.14889627e+01,  -2.86405832e+02],
         [  1.40415805e+02,   3.58726232e+02,   6.47166710e+02]],
         ...

沿dask数组的轴应用函数

例

微信关注