numpy sum函数中如何访问求和迭代索引实现向量化计算

你需要的求和维度迭代索引i不需要通过循环逐次取值，直接用numpy生成对应长度的索引数组，配合广播机制就能完成向量化计算，甚至可以把外层的k循环也一并向量化，进一步压缩运行耗时。

单步替换方案（仅替换内层循环） #

直接用np.arange(len(x))生成从0到len(x)-1的整数序列，这个序列的每个元素就对应你注释代码里的迭代值i，替换你原来写死的常量1即可，修改后的内层计算逻辑和你注释的原生for循环结果完全一致：

import numpy as np
from timeit import default_timer as timer

def get_a(x):
    start = timer()
    n = len(x)
    A = np.zeros(n, dtype=np.float64)
    i_arr = np.arange(n)  # 对应内层循环的所有i值
    for k in range(n):
        # 广播计算所有i对应项，求和等价于内层for循环逻辑
        Ak = np.sum(x * np.exp(-2j * np.pi * (k * i_arr / n)))
        A[k] = np.abs(Ak)  # 复数取模等价于norm计算
    end = timer()
    print(f"单层向量化耗时: {end - start}s")
    return A

这个版本的内层计算完全由numpy的C底层实现执行，速度比纯Python两层循环快10~100倍。

全向量化优化方案（移除所有Python层循环） #

如果要进一步提速，可以把外层循环的k也生成为数组，通过二维广播一次性计算所有k、i组合对应的项，按维度求和直接得到全部结果，完全去掉Python层的for循环：

def get_a_fast(x):
    start = timer()
    n = len(x)
    i_arr = np.arange(n).reshape(1, -1) # 形状调整为(1, n)用于广播
    k_arr = np.arange(n).reshape(-1, 1) # 形状调整为(n, 1)用于广播
    # 广播得到n*n的计算矩阵，按行求和得到所有k对应的Ak值
    Ak_arr = np.sum(x * np.exp(-2j * np.pi * (k_arr * i_arr / n)), axis=1)
    A = np.abs(Ak_arr)
    end = timer()
    print(f"全向量化耗时: {end - start}s")
    return A

高性能生产级方案 #

你当前手写的计算逻辑本质是离散傅里叶变换（DFT），上面两种手写向量化方案的时间复杂度都是O(n²)，当序列长度n较大时依然存在性能瓶颈。直接调用numpy内置的快速傅里叶变换接口即可，底层是优化过的FFT实现，时间复杂度为O(nlogn)，速度比手写向量化版本还要快几个量级，计算结果和手写DFT完全一致：

def get_a_fft(x):
    start = timer()
    A = np.abs(np.fft.fft(x))
    end = timer()
    print(f"FFT版本耗时: {end - start}s")
    return A

你可以自行构造测试数据比对三个版本的输出，数值误差在浮点计算精度范围内完全相等。

内容的提问来源于stack exchange，提问作者Nicolas Avendaño