计算dask数组块中元素的总和

时间:2016-10-17 18:00:02

标签: python dask

我想将函数应用于每个块并返回单个元素,例如,从10x10矩阵我想要对每个2x2块求和。

我尝试了下面看到的一些组合,但我总是得到IndexError

m = da.from_array(np.ones((10,10)), chunks=(2,2))
def compute_block_sum(block):
    return np.array([np.sum(block)])
m.map_blocks(compute_block_sum, chunks=(1,1)).compute()

1 个答案:

答案 0 :(得分:1)

使用默认设置map_blocks假定用户提供的函数返回与输入具有相同维数的numpy数组。因此,您可以通过使用compute_block_sum的numpy切片在None函数中添加第二个空维度,从而使上面的示例工作。

In [1]: import dask.array as da
In [2]: import numpy as np
In [3]: m = da.from_array(np.ones((10,10)), chunks=(2,2))

In [4]: def compute_block_sum(block):
   ...:     return np.array([np.sum(block)])[:, None]

In [5]: m.map_blocks(compute_block_sum, chunks=(1, 1)).compute()
Out[5]: 
array([[ 4.,  4.,  4.,  4.,  4.],
       [ 4.,  4.,  4.,  4.,  4.],
       [ 4.,  4.,  4.,  4.,  4.],
       [ 4.,  4.,  4.,  4.,  4.],
       [ 4.,  4.,  4.,  4.,  4.]])