对“稀疏”numpy数组部分的高效数学操作

Question

我在为我的博士论文做一个模拟时遇到了一个挑战：

我需要优化以下这段代码：

repelling_forces = repelling_force_prefactor * np.exp(-(height_r_t/potential_steepness))

在这段代码中，'height_r_t' 是一个真实的 Numpy 数组，而 'potential_steepness' 是一个标量。'repelling_force_prefactor' 也是一个 Numpy 数组，大部分地方是零，但在一些预先计算的位置上是1，这些位置在运行时不会改变（也就是一个掩码）。

显然，这段代码效率不高，因为只在 'repelling_force_prefactor' 不为零的位置计算指数函数会更合理。

我的问题是，怎么才能以最有效的方式做到这一点呢？

到目前为止，我唯一想到的办法是用 'repelling_force_prefactor' 来定义 'height_r_t' 的切片，然后对这些切片应用 'np.exp'。不过，我发现切片的速度比较慢（不确定这是否普遍适用），而且这个解决方案看起来有点笨拙。

顺便提一下，'repelling_force_prefactor' 中1和0的比例大约是1/1000，而我是在循环中运行这个，所以效率非常重要。

（评论：我不介意使用 Cython，因为我反正到时候需要/想要学习它……但我还是个新手，所以需要一个好的指引/解释。）

Answer 1

切片操作通常比计算所有的指数要快得多。与其直接使用掩码 repelling_force_prefactor 来进行切片，我建议先计算出那些非零的索引，然后用这些索引来进行切片：

# before the loop
indices = np.nonzero(repelling_force_prefactor)

# inside the loop
repelling_forces = np.exp(-(height_r_t[indices]/potential_steepness))

这样一来，repelling_forces 就只会包含非零的结果。如果你需要用这些值来更新一个和 height_r_t 原始形状相同的数组，你可以再次使用索引进行切片，或者使用 np.put() 或类似的函数。

在这种情况下，使用索引列表进行切片会比使用布尔掩码更高效，因为索引列表的长度要短一千倍。至于实际测量性能，当然就看你自己了。

Answer 2

掩码数组正好是为了满足你的需求而设计的。

它的性能和Sven的回答是一样的：

height_r_t = np.ma.masked_where(repelling_force_prefactor == 0, height_r_t)
repelling_forces = np.ma.exp(-(height_r_t/potential_steepness))

掩码数组的好处在于，你不需要去切割和扩展你的数组，数组的大小始终保持不变。不过，numpy会自动知道在掩码位置不进行计算。

另外，你可以对带有不同掩码的数组进行求和，得到的结果数组会是这些掩码的交集。