如何用Numpy替代不同向量场景下的for循环以加速代码

原循环代码 #

用户原始代码通过for循环遍历p的每个元素重复计算，在N=1e6的场景下运行效率极低：

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

用户尝试的向量化代码 #

用户尝试将变量转为矩阵形式，但使用np.outer和np.matrix导致代码冗余，还可能因创建大矩阵降低内存效率：

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))

优化后的向量化实现 #

利用Numpy的广播机制，无需显式创建大尺寸矩阵，同时简化计算逻辑，大幅提升运行速度和内存效率：

import numpy as np

N = 1000000
d = 2000

# 调整形状触发广播：p为行向量(1,211)，alpha为列向量(N,1)
p = np.linspace(0, 210, 211)[np.newaxis, :]
alpha = np.linspace(1e-15, np.pi/2, N)[:, np.newaxis]

# 提前计算三角函数，避免重复运算
cos_alpha = np.cos(alpha)
tan_alpha = np.tan(alpha)

# 简化p1+d1的计算：p1+d1 = (p+d)*cosα
p_plus_d = p + d
p1_plus_d1 = p_plus_d * cos_alpha

# 计算k
k = 1 / (p_plus_d * tan_alpha)

# 利用指数性质简化：exp(A)**B = exp(A*B)
z = np.exp(p1_plus_d1 * k)

优化说明 #

1. 广播机制：通过调整p和alpha的形状，让Numpy自动在(N,211)维度上完成元素级运算，无需手动创建np.outer矩阵，减少内存占用。
2. 计算简化：将p1+d1合并为(p+d)*cosα减少一次数组加法；利用指数运算数学性质减少一次指数计算。
3. 弃用旧API：放弃np.matrix，使用更灵活高效的ndarray，符合Numpy现代用法。

内容的提问来源于stack exchange，提问作者nico.wagner