Python dask.array.map_blocks用法及代码示例

用法:
dask.array.map_blocks(func, *args, name=None, token=None, dtype=None, chunks=None, drop_axis=[], new_axis=None, meta=None, **kwargs)

跨 dask 数组的所有块映射函数。

请注意，map_blocks 将尝试通过在输入的 0-d 版本上调用 func 来自动确定输出数组类型。如果您希望函数在 0-d 数组上操作时不会成功，请参考下面的 meta 关键字参数。

参数：

func：可调用的: 应用于数组中每个块的函数。如果 func 接受 block_info= 或 block_id= 作为关键字参数，则这些将被传递字典，其中包含有关计算期间输入和输出块/数组的信息。有关详细信息，请参阅示例。
args：dask 数组或其他对象
dtype：np.dtype，可选: 输出数组的dtype。建议提供这个。如果未提供，将通过将该函数应用于一小组假数据来推断。
chunks：元组，可选: 如果函数不保留形状，则生成块的块形状。如果未提供，则假定结果数组与第一个输入数组具有相同的块结构。
drop_axis：数字或可迭代，可选: 函数丢失的尺寸。
new_axis：数字或可迭代，可选: 函数创建的新维度。请注意，这些在drop_axis(如果存在)之后应用。
token：字符串，可选: 用于输出数组的键前缀。如果未提供，将根据函数名称确定。
name：字符串，可选: 用于输出数组的键名。请注意，这完全指定了输出键名称，并且必须是唯一的。如果未提供，将由参数的哈希确定。
meta：array-like，可选: 输出数组的meta，当指定时，预期是与在此函数返回的数组上调用.compute() 时返回的类型和dtype 相同的数组。如果未提供，meta 将通过将该函数应用于一小组假数据(通常是 0 维数组)来推断。重要的是要确保func 能够成功完成计算而不会在将 0-d 传递给它时引发异常，否则将需要 meta。如果事先知道输出类型(例如，np.ndarray , cupy.ndarray)，则可以传递此类 dtype 类型的空数组，例如：meta=np.array((), dtype=np.int32)。
**kwargs：: 要传递给函数的其他关键字参数。值必须是常量(不是 dask.arrays)

例子：

>>> import dask.array as da
>>> x = da.arange(6, chunks=3)

>>> x.map_blocks(lambda x: x * 2).compute()
array([ 0,  2,  4,  6,  8, 10])

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

>>> d = da.arange(5, chunks=2)
>>> e = da.arange(5, chunks=2)

>>> f = da.map_blocks(lambda a, b: a + b**2, d, e)
>>> f.compute()
array([ 0,  2,  6, 12, 20])

如果函数改变了块的形状，那么你必须明确地提供块。

>>> y = x.map_blocks(lambda x: x[::2], chunks=((2, 2),))

在指定块方面你有一点自由。如果所有输出块大小都相同，则可以仅提供该块大小作为单个元组。

>>> a = da.arange(18, chunks=(6,))
>>> b = a.map_blocks(lambda x: x[:3], chunks=(3,))

如果函数更改块的维度，您必须指定创建或销毁的维度。

>>> b = a.map_blocks(lambda x: x[None, :, None], chunks=(1, 6, 1),
...                  new_axis=[0, 2])

如果指定了chunks 但未指定new_axis，则推断在左侧添加必要的轴数。

Map_blocks 按块位置对齐块而不考虑形状。在下面的示例中，我们有两个具有相同块数但形状和块大小不同的数组。

>>> x = da.arange(1000, chunks=(100,))
>>> y = da.arange(100, chunks=(10,))

要匹配的相关属性是 numblocks。

>>> x.numblocks
(10,)
>>> y.numblocks
(10,)

如果这些匹配(直到广播规则)，那么我们可以跨块映射任意函数

>>> def func(a, b):
...     return np.array([a.max(), b.max()])

>>> da.map_blocks(func, x, y, chunks=(2,), dtype='i8')
dask.array<func, shape=(20,), dtype=int64, chunksize=(2,), chunktype=numpy.ndarray>

>>> _.compute()
array([ 99,   9, 199,  19, 299,  29, 399,  39, 499,  49, 599,  59, 699,
        69, 799,  79, 899,  89, 999,  99])

您的块函数可以通过接受特殊的 block_info 或 block_id 关键字参数来获取有关它在数组中的位置的信息。在计算期间，它们将包含与 func 的每次调用相关的每个输入和输出块(和 dask 数组)的信息。

>>> def func(block_info=None):
...     pass

这将收到以下信息：

>>> block_info  
{0: {'shape': (1000,),
     'num-chunks': (10,),
     'chunk-location': (4,),
     'array-location': [(400, 500)]},
 None: {'shape': (1000,),
        'num-chunks': (10,),
        'chunk-location': (4,),
        'array-location': [(400, 500)],
        'chunk-shape': (100,),
        'dtype': dtype('float64')}}

block_info 字典的键指示哪个是输入和输出 Dask 数组：

输入 Dask 数组： block_info[0]指第一个输入 Dask 数组。字典键是0因为那是对应于第一个输入 Dask 数组的参数索引。如果多个 Dask 数组已作为输入传递给函数，您可以使用与输入参数对应的数字来访问它们，例如：block_info[1],block_info[2]等(请注意，如果您将多个 Dask 数组作为输入传递给map_blocks，则这些数组必须通过匹配数量的块相互匹配，以及对应的维度，直至广播规则。)
输出 Dask 数组： block_info[None]引用输出 Dask 数组，并包含有关输出块的信息。输出块形状和 dtype 可能与输入块不同。

对于每个 dask 数组，block_info 说明：

shape ：完整的 Dask 数组的形状，
num-chunks ：每个维度中完整数组的块数，
chunk-location ：块位置(例如第一维中的第四块)，以及
array-location ：完整 Dask 数组中的数组位置(例如对应于 40:50 的切片)。

除此之外，block_info 为输出数组(在 block_info[None] 中)说明了两个额外参数：

chunk-shape ：输出块形状，以及
dtype ：输出数据类型。

这些特性可以组合起来从头开始合成一个数组，例如：

>>> def func(block_info=None):
...     loc = block_info[None]['array-location'][0]
...     return np.arange(loc[0], loc[1])

>>> da.map_blocks(func, chunks=((4, 4),), dtype=np.float_)
dask.array<func, shape=(8,), dtype=float64, chunksize=(4,), chunktype=numpy.ndarray>

>>> _.compute()
array([0, 1, 2, 3, 4, 5, 6, 7])

block_id 类似于 block_info 但仅包含 chunk_location ：

>>> def func(block_id=None):
...     pass

这将收到以下信息：

>>> block_id  
(4, 3)

您可以使用可选的token 关键字参数在图中指定生成任务的键名前缀。

>>> x.map_blocks(lambda x: x + 1, name='increment')
dask.array<increment, shape=(1000,), dtype=int64, chunksize=(100,), chunktype=numpy.ndarray>

对于可能无法处理 0-d 数组的函数，还可以使用与预期结果类型匹配的空数组指定 meta。在下面的示例中，在计算 meta 时，func 将导致 IndexError：

>>> da.map_blocks(lambda x: x[2], da.random.random(5), meta=np.array(()))
dask.array<lambda, shape=(5,), dtype=float64, chunksize=(5,), chunktype=numpy.ndarray>

同样，可以为 meta 指定一个非 NumPy 数组，并提供一个 dtype ：

>>> import cupy  
>>> rs = da.random.RandomState(RandomState=cupy.random.RandomState)  
>>> dt = np.float32
>>> da.map_blocks(lambda x: x[2], rs.random(5, dtype=dt), meta=cupy.array((), dtype=dt))  
dask.array<lambda, shape=(5,), dtype=float32, chunksize=(5,), chunktype=cupy.ndarray>

相关用法

注：本文由纯净天空筛选整理自dask.org大神的英文原创作品 dask.array.map_blocks。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。