基于行掩码的二维数组多条件向量化过滤方案需求

需要实现一种向量化解决方案，通过带0掩码的二维数组ma按行过滤二维数组ds。规则为：ma中的0表示ds对应位置可匹配任意数值（掩码作用），最终要找出ds中与ma每一行（非0部分）完全匹配的行，需完成对ma所有行的匹配。

循环实现示例 #

import numpy as np
ma = np.array([ [0, 0,  2,  0,  0,  0,  3,  0],
                [0, 0,  3,  0,  0,  0,  0,  2],
                [0, 0,  2,  1,  0,  0,  0,  0],
                [1, 0,  0,  0,  0,  3,  0,  0]])

ds = np.array([[2, 3, 3, 2, 1, 1, 1, 2],
               [3, 3, 2, 2, 3, 2, 3, 3],
               [3, 3, 2, 2, 3, 3, 3, 2],
               [2, 1, 1, 3, 3, 3, 1, 2],
               [1, 3, 2, 1, 1, 2, 1, 1],
               [2, 3, 3, 2, 1, 1, 3, 2],
               [3, 1, 2, 3, 3, 2, 3, 3],
               [2, 1, 1, 2, 1, 2, 1, 1],
               [2, 3, 3, 1, 3, 2, 3, 3],
               [2, 1, 1, 3, 3, 3, 1, 2]])

result = np.zeros([ds.shape[0], ma.shape[0]], dtype = bool)

for i in range (ma.shape[0]):
    ds_filtered     = np.take(ds,  ma[i,:].nonzero(), axis = 1).squeeze()
    result [:,i]  = np.equal(ds_filtered, ma[i,ma[i,:] != 0] ).all(axis = 1)

预期输出结果 #

array([[False,  True, False, False],
       [ True, False, False, False],
       [ True, False, False, False],
       [False, False, False, False],
       [False, False,  True, False],
       [False,  True, False, False],
       [ True, False, False, False],
       [False, False, False, False],
       [False, False, False, False],
       [False, False, False, False]])

任务列表 #

• 任务1（基础）：改写上述代码，移除FOR循环，实现纯向量化解决方案；
• 任务2（优化）：将任务1的向量化方案适配到CuPy，实现GPU加速；
• 任务3（进阶）：编写CUDA自定义内核，可在CuPy或Numba环境中调用；
• 任务4（高级）：优化任务2或3的方案，使用位数据（np.packbits & np.view）替代布尔数组（单个布尔元素占8位），将数组体积压缩至原来的1/8（应对超400GB的实际数据量）。

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