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。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。