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

Question

我想对每个块和 return 单个元素应用一个函数，例如，从 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()

Answer 1

使用默认设置 map_blocks 假定用户提供的函数 returns 是一个与输入维数相同的 numpy 数组。因此，您可以通过使用带有 None.

的 numpy 切片向 compute_block_sum 函数添加第二个空维度来使上面的示例起作用

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

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

Compute sum of the elements in a chunk of a dask array

python

dask