如何使用 numpy 而不是具有不同向量的 for 循环

Question

我想改进我的代码以使其更快，现在，我有一个 for 循环，我不知道如何用 numpy 函数替换它。

import numpy as np

N = 1000000
d = 2000

p = np.linspace(0,210,211)
alpha = np.linspace(0.00000000000001, np.pi/2, N)
d1 = d*np.cos(alpha)

for i in range(len(p)):
    p1 = p[i]*np.cos(alpha)
    k = 1/((p[i]+d)*np.tan(alpha))
    z = np.exp(p1+d1)**k

首先，我尝试将 p1、d1 和 k 向量化为大小合适的矩阵，但我不知道如何在没有循环的情况下计算 z。 此外，我认为这不是一种有效的方法。

import numpy as np

N = 1000000
d = 2000

p = np.linspace(0,210,211)
alpha = np.linspace(0.00000000000001, np.pi/2, N)
d1 = d*np.cos(alpha)


p1 = np.outer(np.cos(alpha),p)
d1 = np.matrix(d1).T * np.matrix(np.ones(len(p)))
k = 1/(np.outer(np.tan(alpha),p)+np.outer(np.tan(alpha),d))

Answer 1

如果您希望p中每个元素一行，而alpha中每个元素一列，则只需向p添加一个轴，使其成为列向量。 Numpy 的广播负责处理 rest：

import numpy as np

N = 100 # modified to run quickly
d = 20

# reshape p to a column vector
p = np.linspace(0,210,211).reshape((-1, 1))

alpha = np.linspace(0.00000000000001, np.pi/2, N)
d1 = d*np.cos(alpha)

p1 = p*np.cos(alpha)        # shape (211, 100)
k = 1/((p+d)*np.tan(alpha)) 
z = np.exp(p1+d1)**k

在线试用

请注意，幂运算溢出到无穷大，但这与 numpy 无关。

另请注意，虽然此答案确实向您展示了如何矢量化您的操作，但这样做可能没有意义，因为您只保存了 211 次迭代循环。 向量化更大的循环更有意义，例如如果p.size比alpha.size

Answer 2

虽然我认为 Pranav Hosangadi 的回答是正确的并且展示了使用 numpy 的现代方式，但我怀疑它是否是最快的，至少对于给定的参数而言。

查看您的原始代码，我看到的缺陷是：

• 使用for i in range(len(p))而不是for _p in p会更像 pythonic，但不会影响性能。
• 在每次迭代中计算np.cos(alpha)和np.tan(alpha)并不是最优的。 它们应该在循环运行之前计算。
• 你使用了一个循环，它可以被矢量化，但这究竟有多糟糕？

如果固定第二个，去掉for循环后的计算过程主要有两个区别：

1. 您失去了 for 循环本身的开销，它是 p 中的恒定时间查找，以及一些 Python 内部结构。 对于 211 次迭代，您可能仍处于较低的微秒区域。 所以赢不了多少。
2. 要处理的对象变大。 Pranav 的回答缩小了它们的比例，因此这不会对他的解决方案产生太大影响。 但是您的原始参数不同：211 * 1000000 乘以浮点大小意味着我们在 1-2 GiB 的范围内。 这意味着在您的计算过程中，您的 CPU 不会有任何缓存命中并且需要从 RAM 加载所有内容，这比 L1/2/3 缓存具有更高的延迟。 在 for 循环的情况下，大小可能正好合适，这会极大地提高速度。

所以我认为这取决于你的 arrays 的大小和你所在的机器，完全矢量化是慢还是快。 在我当前的机器上，以下是 Pranav 解决方案的 2.5 倍：

cosalpha = np.cos(alpha)
tanalpha = np.tan(alpha)
z = np.zeros((p.size, alpha.size))
for i, _p in enumerate(p):
    p1 = _p*cosalpha
    k = 1/((_p+d)*tanalpha)
    z[i] = np.exp(p1+d1)**k