如何用掩码限制两个NumPy数组间的广播除法运算？

现有两个NumPy数组data和divide_by，原本可通过reshape实现对data的每一行（axis 0）执行广播除法。现在需要仅在与data同形状的mask数组中值为True的位置执行该除法，但直接使用data[mask] /= divide_by会触发形状不匹配的ValueError，求可行实现方法。

相关数组定义如下：

import numpy as np

data = np.array([
    [[10, 10, 10],
     [10, 10, 10],
     [10, 10, 10]],

    [[20, 20, 20],
     [20, 20, 20],
     [20, 20, 20]],

    [[30, 30, 30],
     [30, 30, 30],
     [30, 30, 30]],
], dtype=np.float64)

divide_by = np.array([
    [[10, 10, 1]],
    [[1, 10, 10]],
    [[1, 1, 1]],
], dtype=np.float64)

mask = np.array([
    [[False, True, False],
     [False, False, False],
     [True, False, False]],

    [[True, True, True],
     [False, False, True],
     [False, False, False]],

    [[True, False, False],
     [False, False, False],
     [False, False, False]],
])

预期输出：

[[[10.  1. 10.]
  [10. 10.  1.]
  [10. 10. 10.]]

 [[ 2.  2. 20.]
  [20. 20.  2.]
  [20. 20. 20.]]

 [[ 3. 30. 30.]
  [30. 30. 30.]
  [30. 30. 30.]]]

解决方案 #

直接对divide_by使用相同的掩码提取对应位置的值，让除数和被除数的形状完全匹配，即可完成除法操作：

# 核心执行代码
data[mask] /= divide_by[mask]

# 验证输出
print(data)

原理说明 #

data[mask]会将掩码为True的元素提取为一维数组，若直接用原三维的divide_by做除法，会因形状不匹配报错。而divide_by[mask]同样会提取对应掩码位置的元素，得到和data[mask]形状一致的一维数组，此时就能顺利执行广播除法，最终得到预期结果。

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