重叠计算

一些数组操作需要相邻块之间的边界通信。示例操作包括以下内容:

  • 在图像上卷积一个滤波器

  • 滑动和/均值/最大值,…

  • 搜索图像中可能跨越块边界的类似高斯斑点的图像基元

  • 计算偏导数

  • 玩生命游戏_Life

Dask Array 通过创建一个新数组来支持这些操作,其中每个块都由其邻居的边界稍微扩展。这会产生额外的复制成本和许多小块的通信,但允许局部函数以令人尴尬的并行方式进行评估。

这些计算的主要API是下面定义的 map_overlap 方法:

map_overlap(func, *args[, depth, boundary, ...])

在具有一定重叠的数组块上应用函数

dask.array.map_overlap(func, *args, depth=None, boundary=None, trim=True, align_arrays=True, allow_rechunk=True, **kwargs)[源代码]

在具有一定重叠的数组块上应用函数

我们在数组的块之间共享相邻区域,映射一个函数,然后修剪掉相邻的条带。如果深度大于特定轴上的任何块,则数组会被重新分块。

请注意,此函数将在计算之前尝试自动确定输出数组类型,如果您预计该函数在操作0-d数组时不会成功,请参阅``map_blocks``中的``meta``关键字参数。

参数
func: 函数

应用于每个扩展块的函数。如果提供了多个数组,那么函数应预期按相同顺序接收每个数组的块。

参数dask 数组
depth: int, tuple, dict 或 list, 仅关键字参数

每个块应与其邻居共享的元素数量。如果是一个元组或字典,则每个轴可以不同。如果是一个列表,则该列表的每个元素必须是定义相应数组深度的整数、元组或字典。可以使用 (-/+) 元组的字典值指定不对称深度。请注意,当 boundary 为 ‘none’ 时,目前仅支持不对称深度。默认值为 0。

boundary: str, tuple, dict 或 list, 仅关键字参数

如何处理边界。值可以是 ‘reflect’、’periodic’、’nearest’、’none’,或任何常数值,如 0 或 np.nan。如果是一个列表,则每个元素必须是一个字符串、元组或字典,定义 args 中相应数组的边界。默认值是 ‘reflect’。

trim: bool, 仅关键字

是否在调用映射函数后从每个块中修剪 depth 元素。如果你的映射函数已经为你完成了这项工作,请将其设置为 False。

align_arrays: bool, 仅关键字参数

当提供多个数组时,是否沿等尺寸的维度对齐块。这允许某些数组中的较大块被分解为与其它数组中的块大小匹配的较小块,以便它们能够兼容块函数映射。如果这是false,那么如果数组在每个维度中没有相同数量的块,则会抛出错误。

allow_rechunk: bool, 仅关键字参数

允许重新分块,否则块大小需要匹配,核心维度只能由一个块组成。

**kwargs:

map_blocks 中有效的其他关键字参数

示例

>>> import numpy as np
>>> import dask.array as da
>>> x = np.array([1, 1, 2, 3, 3, 3, 2, 1, 1])
>>> x = da.from_array(x, chunks=5)
>>> def derivative(x):
...     return x - np.roll(x, 1)
>>> y = x.map_overlap(derivative, depth=1, boundary=0)
>>> y.compute()
array([ 1,  0,  1,  1,  0,  0, -1, -1,  0])
>>> x = np.arange(16).reshape((4, 4))
>>> d = da.from_array(x, chunks=(2, 2))
>>> d.map_overlap(lambda x: x + x.size, depth=1, boundary='reflect').compute()
array([[16, 17, 18, 19],
       [20, 21, 22, 23],
       [24, 25, 26, 27],
       [28, 29, 30, 31]])
>>> func = lambda x: x + x.size
>>> depth = {0: 1, 1: 1}
>>> boundary = {0: 'reflect', 1: 'none'}
>>> d.map_overlap(func, depth, boundary).compute()  
array([[12,  13,  14,  15],
       [16,  17,  18,  19],
       [20,  21,  22,  23],
       [24,  25,  26,  27]])

da.map_overlap 函数也可以接受多个数组。

>>> func = lambda x, y: x + y
>>> x = da.arange(8).reshape(2, 4).rechunk((1, 2))
>>> y = da.arange(4).rechunk(2)
>>> da.map_overlap(func, x, y, depth=1, boundary='reflect').compute() 
array([[ 0,  2,  4,  6],
       [ 4,  6,  8,  10]])

当给出多个数组时,它们不需要具有相同的维度,但它们必须能够一起广播。数组按块对齐(就像在 da.map_blocks 中一样),因此块必须具有共同的块大小。只要 align_arrays 为 True,这种共同的块划分会自动确定。

>>> x = da.arange(8, chunks=4)
>>> y = da.arange(8, chunks=2)
>>> r = da.map_overlap(func, x, y, depth=1, boundary='reflect', align_arrays=True)
>>> len(r.to_delayed())
4
>>> da.map_overlap(func, x, y, depth=1, boundary='reflect', align_arrays=False).compute()
Traceback (most recent call last):
    ...
ValueError: Shapes do not align {'.0': {2, 4}}

另请注意,此函数默认等同于 map_blocks。 对于提供给 func 的数组中出现的任何重叠,必须定义非零的 depth

>>> func = lambda x: x.sum()
>>> x = da.ones(10, dtype='int')
>>> block_args = dict(chunks=(), drop_axis=0)
>>> da.map_blocks(func, x, **block_args).compute()
np.int64(10)
>>> da.map_overlap(func, x, **block_args, boundary='reflect').compute()
np.int64(10)
>>> da.map_overlap(func, x, **block_args, depth=1, boundary='reflect').compute()
np.int64(12)

对于可能不处理 0-d 数组的函数,也可以使用与预期结果类型匹配的空数组来指定 meta。在下面的示例中,计算 metafunc 将导致 IndexError

>>> x = np.arange(16).reshape((4, 4))
>>> d = da.from_array(x, chunks=(2, 2))
>>> y = d.map_overlap(lambda x: x + x[2], depth=1, boundary='reflect', meta=np.array(()))
>>> y
dask.array<_trim, shape=(4, 4), dtype=float64, chunksize=(2, 2), chunktype=numpy.ndarray>
>>> y.compute()
array([[ 4,  6,  8, 10],
       [ 8, 10, 12, 14],
       [20, 22, 24, 26],
       [24, 26, 28, 30]])

同样地,可以为 meta 指定一个非 NumPy 数组:

>>> import cupy  
>>> x = cupy.arange(16).reshape((4, 4))  
>>> d = da.from_array(x, chunks=(2, 2))  
>>> y = d.map_overlap(lambda x: x + x[2], depth=1, boundary='reflect', meta=cupy.array(()))  
>>> y  
dask.array<_trim, shape=(4, 4), dtype=float64, chunksize=(2, 2), chunktype=cupy.ndarray>
>>> y.compute()  
array([[ 4,  6,  8, 10],
       [ 8, 10, 12, 14],
       [20, 22, 24, 26],
       [24, 26, 28, 30]])

解释

考虑Dask数组中的两个相邻块:

两个不相交的相邻块。

我们通过在数组之间交换附近的薄切片来扩展每个块:

两个相邻的块,沿着它们的共享边界有细条,表示它们之间的共享数据。

我们在所有方向上进行此操作,包括与重叠函数的对角交互:

一个二维的块网格,每个块的边界周围有细条,代表从邻居共享的数据。它们还包括用于从对角邻居共享数据的小角位。
>>> import dask.array as da
>>> import numpy as np

>>> x = np.arange(64).reshape((8, 8))
>>> d = da.from_array(x, chunks=(4, 4))
>>> d.chunks
((4, 4), (4, 4))

>>> g = da.overlap.overlap(d, depth={0: 2, 1: 1},
...                       boundary={0: 100, 1: 'reflect'})
>>> g.chunks
((8, 8), (6, 6))

>>> np.array(g)
array([[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
       [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
       [  0,   0,   1,   2,   3,   4,   3,   4,   5,   6,   7,   7],
       [  8,   8,   9,  10,  11,  12,  11,  12,  13,  14,  15,  15],
       [ 16,  16,  17,  18,  19,  20,  19,  20,  21,  22,  23,  23],
       [ 24,  24,  25,  26,  27,  28,  27,  28,  29,  30,  31,  31],
       [ 32,  32,  33,  34,  35,  36,  35,  36,  37,  38,  39,  39],
       [ 40,  40,  41,  42,  43,  44,  43,  44,  45,  46,  47,  47],
       [ 16,  16,  17,  18,  19,  20,  19,  20,  21,  22,  23,  23],
       [ 24,  24,  25,  26,  27,  28,  27,  28,  29,  30,  31,  31],
       [ 32,  32,  33,  34,  35,  36,  35,  36,  37,  38,  39,  39],
       [ 40,  40,  41,  42,  43,  44,  43,  44,  45,  46,  47,  47],
       [ 48,  48,  49,  50,  51,  52,  51,  52,  53,  54,  55,  55],
       [ 56,  56,  57,  58,  59,  60,  59,  60,  61,  62,  63,  63],
       [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
       [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]])

边界

关于重叠,您可以指定如何处理边界。当前的策略包括以下内容:

  • periodic - 将边界环绕到另一侧

  • reflect - 向外反射每条边框

  • any-constant - 用此值填充边界

一个示例的边界类型参数可能如下所示:

{0: 'periodic',
 1: 'reflect',
 2: np.nan}

或者,你可以使用 dask.array.pad() 进行其他类型的填充。

在块上应用函数

重叠与跨块映射函数紧密相关。现在,该函数可以使用从邻居复制过来的额外信息,这些信息在每个块中并不本地存储:

>>> from scipy.ndimage import gaussian_filter
>>> def func(block):
...    return gaussian_filter(block, sigma=1)

>>> filt = g.map_blocks(func)

在这种情况下,我们使用了 SciPy 函数,但也可以使用任何其他任意函数。这是与 Numba 的一个良好交互点。

如果你的函数不保留块的形状,那么你需要提供一个 chunks 关键字参数。如果你的块大小是规则的,那么这个参数可以接受一个块形状,例如 (1000, 1000)。在块大小不规则的情况下,它必须是一个包含完整块形状的元组,例如 ((1000, 700, 1000), (200, 300))

>>> g.map_blocks(myfunc, chunks=(5, 5))

如果你的函数需要知道它所操作的块的位置,你可以给你的函数一个关键字参数 block_id:

def func(block, block_id=None):
    ...

这个额外的关键字参数将被赋予一个元组,该元组提供块的位置,例如 (0, 0) 表示左上角的块,或 (0, 1) 表示紧邻该块右侧的块。

修剪多余

在映射一个被阻塞的函数后,你可能希望从每个块中修剪掉它们被扩展的相同数量的边界。函数 trim_internal 在这里很有用,并且接受与 overlap 相同的 depth 参数:

>>> x.chunks
((10, 10, 10, 10), (10, 10, 10, 10))

>>> y = da.overlap.trim_internal(x, {0: 2, 1: 1})
>>> y.chunks
((6, 6, 6, 6), (8, 8, 8, 8))

完整工作流程

因此,一个非常典型的重叠工作流程包括 overlapmap_blockstrim_internal

>>> x = ...
>>> g = da.overlap.overlap(x, depth={0: 2, 1: 2},
...                        boundary={0: 'periodic', 1: 'periodic'})
>>> g2 = g.map_blocks(myfunc)
>>> result = da.overlap.trim_internal(g2, {0: 2, 1: 2})