如何用Dask并行化处理分块3D数组并传递多参数函数？

问题背景 #

我有一个维度为mid_date,y,x的数据集ds，其中包含3D numpy数组ds.v。希望对数组的每个y,x单元应用一个函数，返回维度更低的向量。示例代码如下：

# 加载输入矩阵    
array_input = ds.v.values

# 定义m <= array_input.shape[0]
m = 10

# 定义函数：输入长度为array_input.shape[0]的向量，返回长度为m的向量
def func(pt_in, m):
    # 降维处理
    pt_out = pt_in[:m]
    return pt_out

# 遍历每个y,x单元应用函数，存储结果到输出数组
array_output = np.zeros((m, ds.v.values.shape[1], ds.v.values.shape[2]))
for i in range(ds.v.values.shape[1]):
    for j in range(ds.v.values.shape[2]):
           array_output[:,i,j] = func(array_input[:,i,j], m)

我的目标：

• 并行化函数，遍历输入数组的y, x维度，将结果存储为尺寸为(m, y, x)的输出数组。
• 使用Dask以分块方式应用函数，避免加载整个数据集到内存。

尝试过dask_array.map_blocks，但实际函数有6个输入参数，只有第一个是变化的待处理单元，而map_blocks传递的是块而非单个单元，无法成功传递参数。

解决方案 #

1. 适配函数处理Dask块 #

Dask的map_blocks针对整块数据操作，而非单个y,x单元。我们需要把函数改造成能处理整个块的形式，同时通过args参数传递固定参数：

import dask.array as da
import xarray as xr
import numpy as np

# 加载Zarr数据集为Dask数组（仅加载元数据，不占内存）
ds = xr.open_zarr('http://its-live-data.s3.amazonaws.com/datacubes/v02/N50W140/ITS_LIVE_vel_EPSG3413_G0120_X-3350000_Y350000.zarr')
dask_array = ds.v.data  # 形状：(mid_date, y, x)

# 定义参数
m = 10
# 示例其他固定参数
param1 = 1
param2 = 2
param3 = 3
param4 = 4
param5 = 5

# 适配为处理Dask块的函数
def process_block(block, m, p1, p2, p3, p4, p5):
    # block形状：(mid_date_chunk, y_chunk, x_chunk)
    # 对块内每个y,x单元应用函数，输出形状：(m, y_chunk, x_chunk)
    # 示例用切片实现降维，实际替换为你的函数逻辑
    output = np.zeros((m, block.shape[1], block.shape[2]), dtype=block.dtype)
    for i in range(block.shape[1]):
        for j in range(block.shape[2]):
            output[:, i, j] = func(block[:, i, j], m, p1, p2, p3, p4, p5)
    return output

# 更高效的向量化替代方案（推荐）
def process_block_vectorized(block, m, p1, p2, p3, p4, p5):
    # 沿mid_date轴对每个y,x单元应用函数
    return np.apply_along_axis(
        lambda x: func(x, m, p1, p2, p3, p4, p5),
        axis=0,
        arr=block
    ).transpose(2, 0, 1)  # 将输出从(y,x,m)转为(m,y,x)

# 使用map_blocks处理
output_dask = da.map_blocks(
    process_block_vectorized,
    dask_array,
    args=(m, param1, param2, param3, param4, param5),
    dtype=dask_array.dtype,
    chunks=(m, dask_array.chunks[1], dask_array.chunks[2])  # 匹配输入的y,x分块
)

2. 执行与存储 #

• 按需计算结果（仅加载当前处理的块到内存）：

  output_array = output_dask.compute()
  

• 直接保存为Zarr文件，避免内存过载：

  output_xr = xr.DataArray(
      output_dask,
      dims=['new_dim', 'y', 'x'],
      coords={'y': ds.y, 'x': ds.x}
  )
  output_xr.to_zarr('output.zarr', mode='w')
  

3. 简洁替代：Dask apply_along_axis #

如果函数可沿轴处理，直接用da.apply_along_axis更简洁：

from functools import partial

# 封装固定参数
func_wrapped = partial(func, m=m, param1=param1, param2=param2, param3=param3, param4=param4, param5=param5)

# 沿mid_date轴（轴0）应用函数，转置调整维度
output_dask = da.apply_along_axis(func_wrapped, axis=0, arr=dask_array).transpose(2, 0, 1)

内容的提问来源于stack exchange，提问作者Nihilum