如何用低维布尔数组切片Dask数组？复现Numpy掩码切片行为

Dask数组的掩码索引机制不支持Numpy那样的跨维度自动广播处理，直接用2D掩码切片3D数组会抛出IndexError。要复现Numpy的行为，可以通过将掩码转换为线性索引，结合数组重塑和da.take操作实现，具体步骤如下：

解决方法 #

1. 计算掩码在展平后的后两维中的线性索引；
2. 将3D Dask数组重塑为(第一维长度, 后两维总元素数)的2D数组；
3. 使用da.take按索引提取目标元素。

完整代码示例 #

import numpy as np
import dask.array as da

# 创建随机值3D数组，以及与第二、三维形状匹配的掩码
y_da = da.random.random(size=(20, 100, 100))
mask = np.zeros((100, 100), dtype=np.uint8)
mask[20:80, 3:77] = 1

# 计算掩码的线性索引（展平后两维）
flat_mask_indices = np.where(mask.flatten())[0]

# 重塑Dask数组，将后两维合并为一维
y_da_reshaped = y_da.reshape(20, -1)

# 按索引提取元素
y_da_sliced = y_da_reshaped.take(flat_mask_indices, axis=1)

# 验证结果（与Numpy输出一致）
print(y_da_sliced.compute().shape)  # 输出 (20, 4440)

原理说明 #

Numpy执行y_np[:, mask == 1]时，会自动将数组后两维展平，再根据掩码筛选元素。但Dask的分区存储模型无法直接处理这种跨维度的不规则掩码——它需要明确的分区索引范围来调度计算。通过将掩码转换为固定线性索引，我们把不规则掩码操作转化为Dask支持的take操作，从而实现和Numpy一致的切片效果。

内容的提问来源于stack exchange，提问作者Loïc Dutrieux