如何用dask.map_blocks和scipy.zoom缩放Dask数组并保存为Zarr/HDF5

问题场景 #

使用SciPy的zoom函数对大型Dask张量数组进行缩放后，尝试通过dask.array.to_hdf5或dask.array.to_zarr保存时，出现以下错误：

TypeError: Can't broadcast (40, 40) -> (20, 20)

示例代码：

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

data = da.random.randint(0,100, (100,100), chunks=(20,20))
data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16)
data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

错误原因 #

da.map_blocks默认沿用原数组的分块大小（20x20），但zoom(x,2)将每个输入块缩放为40x40，导致保存时HDF5/Zarr期望的分块尺寸与实际输出块尺寸不匹配，触发广播错误。

解决方法 #

方法1：显式指定输出分块大小（推荐） #

在map_blocks中通过chunks参数直接定义缩放后的分块尺寸，让Dask明确每个输出块的大小：

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

data = da.random.randint(0,100, (100,100), chunks=(20,20))
# 每个20x20块缩放后变为40x40，显式指定chunks参数
data_upsampled = da.map_blocks(
    lambda x: zoom(x, 2),
    data,
    dtype=np.uint16,
    chunks=(40, 40)
)
data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

方法2：缩放后重新调整分块 #

如果需要自定义输出分块大小，可先完成缩放，再用da.rechunk调整：

data_upsampled = da.map_blocks(lambda x: zoom(x,2), data, dtype = np.uint16)
# 调整为目标分块大小，例如(50,50)
data_upsampled = data_upsampled.rechunk((50, 50))
data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

方法3：使用Dask原生缩放函数（更高效） #

若无需SciPy zoom的特定插值算法，可使用Dask原生函数自动处理分块：

• da.resize：直接指定目标尺寸

data_upsampled = da.resize(data, (200, 200))
data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

• da.coarsen（上采样版本）：适合整数倍缩放，通过重复元素实现

# 对每个维度重复2次，实现2倍上采样
data_upsampled = da.coarsen(
    lambda x: np.repeat(np.repeat(x, 2, axis=0), 2, axis=1),
    data,
    {0: 2, 1: 2}
)
data_upsampled.to_hdf5('myfile.hdf5', '/up_sampled')

注意：如果必须使用SciPy zoom（如需要双线性/三次插值），优先选择方法1，避免额外的分块调整开销。

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