求助：Dask map_overlap处理缩放后数组的分块不匹配报错解决

需求说明 #

我有两个类矩阵变量：维度(100,100)的a和维度(200,200)的b。需要先将a放大2倍以匹配b的尺寸，再通过Dask的map_overlap函数执行自定义计算并指定重叠区域，最终得到与b尺寸一致的结果。

复现代码 #

import dask.array as da
import numpy as np
from scipy.ndimage import zoom

data = da.ones((40, 140,140), chunks=(10, 40,40))
data2 = da.ones((80, 280, 280), chunks=(10, 40,40))
data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16, chunks=(20,80,80))

res = da.map_overlap(lambda x,y: x+y, data_upsampled, data2, depth = (5,10,10), boundary="constant")

错误信息 #

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[7], line 10
      7 data2 = da.ones((80, 280, 280), chunks=(10, 40,40))
      8 data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16, chunks=(20,80,80))
---> 10 res = da.map_overlap(lambda x,y: x+y, data_upsampled, data2, depth = (5,10,10), boundary="constant")
     11 #data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

File \AppData\Local\anaconda3\envs\napari\lib\site-packages\dask\array\overlap.py:693, in map_overlap(func, depth, boundary, trim, align_arrays, allow_rechunk, *args, **kwargs)
    690 if align_arrays:
    691     # Reverse unification order to allow block broadcasting
    692     inds = [list(reversed(range(x.ndim))) for x in args]
--> 693     _, args = unify_chunks(*list(concat(zip(args, inds))), warn=False)
    695 # Escape to map_blocks if depth is zero (a more efficient computation)
    696 if all([all(depth_val == 0 for depth_val in d.values()) for d in depth]):

File \AppData\Local\anaconda3\envs\napari\lib\site-packages\dask\array\core.py:3971, in unify_chunks(*args, **kwargs)
   3968     else:
   3969         nameinds.append((a, ind))
-> 3971 chunkss = broadcast_dimensions(nameinds, blockdim_dict, consolidate=common_blockdim)
   3972 nparts = math.prod(map(len, chunkss.values()))
   3974 if warn and nparts and nparts >= max_parts * 10:

File \AppData\Local\anaconda3\envs\napari\lib\site-packages\dask\blockwise.py:1467, in broadcast_dimensions(argpairs, numblocks, sentinels, consolidate)
   1464 g2 = {k: v - set(sentinels) if len(v) > 1 else v for k, v in g.items()}
...
   3883 # burned through all of the chunk tuples.
   3884 # For efficiency's sake we reverse the lists so that we can pop off the end
   3885 rchunks = [list(ntd)[::-1] for ntd in non_trivial_dims]

ValueError: ('Chunks do not add up to same value', {(40, 40, 40, 40, 40, 40, 40), (80, 80, 80, 80)})

问题原因 #

map_overlap要求输入数组的块结构（chunks）在对应维度上必须完全兼容：

• data2的第二、三维块划分是(40,40,40,40,40,40,40)（7个40，总和280）
• data_upsampled的第二、三维实际块划分是(80,80,80,40)（与指定的(80,80,80,80)矛盾，因为原维度140放大2倍是280，无法被4个80整除）
  两者的块划分无法对齐，导致报错。

解决方案 #

核心是让两个输入数组的块结构完全一致，有两种实现方式：

方式1：调整data2的块结构匹配data_upsampled #

import dask.array as da
import numpy as np
from scipy.ndimage import zoom

data = da.ones((40, 140,140), chunks=(10, 40,40))
# 调整data2的块结构，与data_upsampled的块划分对齐
data2 = da.ones((80, 280, 280), chunks=(10, [80,80,80,40], [80,80,80,40]))
data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16, chunks=(20,80,80))

res = da.map_overlap(lambda x,y: x+y, data_upsampled, data2, depth=(5,10,10), boundary="constant")

方式2：将data_upsampled重新分块匹配data2 #

import dask.array as da
import numpy as np
from scipy.ndimage import zoom

data = da.ones((40, 140,140), chunks=(10, 40,40))
data2 = da.ones((80, 280, 280), chunks=(10, 40,40))
data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16)
# 重新分块，对齐data2的块结构
data_upsampled = data_upsampled.rechunk(chunks=(10,40,40))

res = da.map_overlap(lambda x,y: x+y, data_upsampled, data2, depth=(5,10,10), boundary="constant")

两种方式都能让数组块结构对齐，确保map_overlap正常执行。

内容的提问来源于stack exchange，提问作者Nim.Moj