沿 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 bits.

从性能角度来看，您可能还需要考虑使用更大的块大小。如果每个块有 6000 个数字，您就有超过一百万个块，这意味着您可能会在调度上花费比实际计算更多的时间。一般来说，我会拍摄几兆字节大小的块。当然，增加块大小会导致映射函数变得更加复杂。

Example

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]],
         ...