Numpy - 以向量形式编写函数？

Question

我对 NumPy（或 SciPy）很陌生并且来自 Octave/Matlab，这对我来说似乎有点挑战。

我正在阅读文档并编写一些基本功能。 我遇到了这一节： 矢量化函数（vectorize）

它定义了这个函数：

def addsubtract(a, b):
   if a > b:
       return a - b
   else:
       return a + b

然后向量化它：

vec_addsubtract = np.vectorize(addsubtract)

但最后，它说：

这个特殊的函数本可以在不使用 vectorize 的情况下以向量形式编写。

我不知道写这样的函数的任何其他方式。 那么什么是向量形式呢？

Answer 1

np.vectorize是一个美化的 python for循环，这意味着它有效地剥离了 numpy 提供的任何优化。

为了实际矢量化addsubtract ，我们可以利用 numpy 提供三样东西的事实：矢量化add函数、矢量化subtract函数和各种布尔掩码操作。

最简单但效率最低的方法是使用np.where ：

np.where(a > b, a - b, a + b)

这是低效的，因为它在所有情况下都预先计算a - b和a + b ，然后为每个元素从一个或另一个中进行选择。

更有效的解决方案只会计算条件需要的值：

result = np.empty_like(a)
mask = a > b
np.subtract(a, b, where=mask, out=result)
np.add(a, b, where=~mask, out=result)

对于非常小的数组，复杂方法的开销使其不值得。 但对于大型阵列，这是最快的解决方案。

有趣的事实：您引用的教程中的页面在 SciPy 教程的未来版本中将不可用，正是因为它是 NumPy 的介绍，如PR #12432 中所述。

Answer 2

您可以使用np.where执行此np.where ，它计算两个结果（ ab和a+b ）并根据布尔数组（ a>b ）选择值：

def addsubtract(a, b):
    return np.where(a>b, a-b, a+b)

它可以看作是一个向量化的三元运算符：“当 a>b 时，从 ab 中取值，否则从 a+b 中取值”。

尽管计算了两种可能的结果，但它比您编写的矢量化 if/else 函数（至少在我的机器上）要快得多。