优化二维余弦函数的Numpy实现：移除循环提升运行速度

优化后的代码 #

import numpy as np

k = np.linspace(0, 4.76 * 10, 2400)
kx, ky = np.meshgrid(k, k)

# 将列表转为Numpy数组，适配后续广播操作
rx = np.array([0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4])
ry = np.array([0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2])

# 扩展维度实现广播，让rx/ry与kx/ky自动匹配计算维度
rx_broadcast = rx[:, np.newaxis, np.newaxis]
ry_broadcast = ry[:, np.newaxis, np.newaxis]

# 向量化计算所有余弦项并求和，直接得到最终z数组
z = 0.05 * np.sum(np.cos(2 * np.pi * (rx_broadcast * kx + ry_broadcast * ky)), axis=0)

核心优化点 #

• 移除Python循环：原代码的for循环在Python解释器层面逐个处理元素，效率极低。改用Numpy向量化操作后，计算逻辑由底层C实现，能充分利用CPU并行能力，当rx/ry长度增加时速度优势会指数级放大。
• 替换冗余张量运算：原tensordot生成的高维数组会占用额外内存，改用np.newaxis扩展维度实现广播，既能完成对应元素的乘法运算，又能大幅降低内存开销。
• 批量求和替代累加：用np.sum直接对所有维度的余弦项求和，避免了循环中反复的数组赋值操作，进一步提升计算效率。

内容的提问来源于stack exchange，提问作者mikanim