python - 参数化数组的矢量化加法

Question

我想对由(复杂的)参数函数定义的一系列向量求和。有没有一种方法可以使用矢量化代替嵌套的 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)

Answer 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) array 它会假装第二个数组也是一个 (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 强大功能的关键，但不可否认，我也花了一些时间来熟悉它。