dask 数组 map_blocks，以不同形状的 dask 数组作为参数

Question

我正在尝试使用 dask.array.map_blocks 来处理一个 dask 数组，使用具有不同形状的第二个 dask 数组作为参数。该用例首先在二维图像堆栈（4 维）上运行一些峰值查找，该图像作为 np.objects 的二维 dask 数组返回。因此，两个 dask 数组的两个第一个维度是相同的。然后使用峰值从 4 维数据集中提取强度。在下面的代码中，我省略了峰值查找部分。 Dask 版本 1.0.0。

import numpy as np
import dask.array as da

def test_processing(data_chunk, position_chunk):
    output_array = np.empty(data_chunk.shape[:-2], dtype='object')
    for index in np.ndindex(data_chunk.shape[:-2]):
        islice = np.s_[index]
        intensity_list = []
        data = data_chunk[islice]
        positions = position_chunk[islice]
        for x, y in positions:
            intensity_list.append(data[x, y])
        output_array[islice] = np.array(intensity_list)
    return output_array

data = da.random.random(size=(4, 4, 10, 10), chunks=(2, 2, 10, 10))
positions = np.empty(data.shape[:-2], dtype='object')
for index in np.ndindex(positions.shape):
    positions[index] = np.arange(10).reshape(5, 2)

data_output = da.map_blocks(test_processing, data, positions, dtype=np.object,
                            chunks=(2, 2), drop_axis=(2, 3))
data_output.compute()

这给出了错误ValueError: Can't drop an axis with more than 1 block. Please useatopinstead.，我猜这是由于positions 有3 个维度，而data 有4 个维度。

相同的功能，但没有positions dask 数组可以正常工作。

import numpy as np
import dask.array as da

def test_processing(data_chunk):
    output_array = np.empty(data_chunk.shape[:-2], dtype='object')
    for index in np.ndindex(data_chunk.shape[:-2]):
        islice = np.s_[index]
        intensity_list = []
        data = data_chunk[islice]
        positions = [[5, 2], [1, 3]]
        for x, y in positions:
            intensity_list.append(data[x, y])
        output_array[islice] = np.array(intensity_list)
    return output_array

data = da.random.random(size=(4, 4, 10, 10), chunks=(2, 2, 10, 10))

data_output = da.map_blocks(test_processing, data, dtype=np.object,
                            chunks=(2, 2), drop_axis=(2, 3))
data_computed = data_output.compute()

Answer 1

这已在 dask 的更新版本中得到修复：在 dask 的 2.3.0 版上运行相同的代码可以正常工作。