8

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

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

这是我获取三次多项式系数的函数（time_axis轴值是在别处定义的全局 1D 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?

4

2 回答 2

8

我怀疑如果您的函数返回一个与它所使用的维度相同的数组，您的体验会更顺畅。例如，您可能会考虑按如下方式定义您的函数：

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

然后你可能会忽略new_axis,drop_axis位。

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

例子

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]],
         ...
于 2016-03-29T04:09:22.797 回答
3

派对有点晚了，但认为这可以使用基于 Dask 新功能的替代答案。特别是，我们添加了apply_along_axis，它的行为基本上类似于NumPyapply_along_axis ，除了 Dask Arrays。这导致语法稍微简单一些。此外，它还避免了在将自定义函数应用于每个 1-D 块之前重新分块数据的需要，并且对初始分块没有真正的要求，它试图在最终结果中保留（除了减少或替换的轴） .

In [1]: import dask.array as da

In [2]: import numpy as np

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

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

In [5]: da.apply_along_axis(f, 0, x).compute()
Out[5]: 
array([[[[  2.13570599e+02,   2.28924503e+00,   6.16369231e+01],
         [  4.32000311e+00,   7.01462518e+01,  -1.62215514e+02],
         [  2.89466687e+02,  -1.35522215e+02,   2.86643721e+02]],
         ...
于 2017-10-19T19:05:25.380 回答