参数化数组的向量化加法答案

【问题标题】：vectorized addition of parameterized arrays参数化数组的向量化加法
【发布时间】：2019-12-16 08:34:10
【问题描述】：

我想对由（复杂的）参数函数定义的一系列向量求和。有没有办法在下面的代码中使用矢量化代替嵌套的 for 循环来为 ans 获得相同的值？内存使用不应成为约束。请注意，对于实际问题，已知嵌套 for 循环会限制性能，尽管此处并非如此。

import numpy as np

a1_vec = np.array([0.3, 1])
a2_vec = np.array([3.3, 10])
b1_vec = np.array([0.5, 0.7])
b2_vec = np.array([1.5, 1.3])

x = np.arange(0, 10000)
ans = 0
for a1, b1 in zip(a1_vec, b1_vec):
    for a2, b2 in zip(a2_vec, b2_vec):
        ans += x*np.exp(- a1 - b2) + x**(1 / 2)*np.cos(b1) + x**(1 / 3)*np.sin(a2)

【问题讨论】：

x[:,None,None] * np.exp(-a1_vec[None,None,:]-b2_vec[None,:,None] 应该给出一个 (10000,2,2) 数组，可以是 np.sum(...)。其他术语类似
@hpaulj 你能说更多关于你在评论中发表的内容吗？如果您添加答案，我很乐意接受。
有趣 - 嵌套循环代码实际上更快。我想知道用 numpy、cython 是否有更快的方法来做到这一点？
这确实很有趣且出乎意料，显然python3实际上加强了它的游戏。这甚至适用于更长的 x 向量吗？所以 1e5 或 1e6 在范围内？
您可以尝试的另一件事是 [numba][numba.pydata.org/] 代码被编译的地方。我会在 cython 之前尝试这个。平心而论，您的问题似乎足够小，为什么还要进一步优化？

标签： python arrays numpy vectorization

【解决方案1】：

所以正如@anon 在他的评论中指出的那样，你可以使用 numpy 的数组broadcasting：使用X[:,None,None] 或更明确的X[:,np.newaxis,np.newaxis]（这只是None 的更明确的别名）。使用此np.newaxis，您可以在数组中引入一个新的空维度。您的X.shape（数组的维数）将不再是(10000)，而是(10000,1,1)。

对两个数组的任何操作都需要在两个数组上具有相同的维度（例如，(2,5) 和 (2,5) 将起作用）。 numpy 的一个很酷的特性是广播，当在两个数组之间进行数组操作时会发生这种情况，其中一个大于 1 的维度与 1 的维度相结合（例如，(2,5) 和 (2,1)，其中第二个维度在第一个向量中为 5，在第二个向量中为 1）。在这种情况下，numpy 将通过简单地循环来广播维度为 1 的数组（并且它将以 c 编译的快速方式这样做）。在示例中，它将使用 (2,5) 数组和 (2,1) 数组，它会假装第二个数组也是 (2,5) 数组，只是相同的两个数字重复了 5 次。

作为一个简单的例子来看看

a = np.arange(3)
print(a*a)
#[0,2,4]

这只是普通的元素乘法

现在如果你使用广播规则引入空维度：

print(a[:,None]*a[None,:])
#[[0 0 0]
# [0 1 2]
# [0 2 4]]

这真的很酷，也是理解 numpy 强大功能的关键，但我承认我也花了一些时间来熟悉它。

【讨论】：