Python实战：用NumPy手写离散傅里叶变换（附完整代码对比FFT性能）

从零实现离散傅里叶变换：深入理解FFT为何如此之快

如果你曾经在信号处理或数据分析中用过NumPy的np.fft.fft函数，可能会觉得它快得不可思议。输入一个数组，几乎瞬间就能得到频域结果。但你是否好奇过，这背后究竟发生了什么？为什么我们自己写的“朴素”傅里叶变换代码会慢上几十甚至几百倍？今天，我们就抛开库函数，亲手用NumPy从最基础的数学公式开始，一步步实现一个完整的离散傅里叶变换，并在这个过程中，彻底弄懂FFT（快速傅里叶变换）算法的精妙之处。这不仅是一次编程练习，更是一次深入理解信号处理核心思想的旅程。

对于数据科学家、音频工程师或任何需要处理时序数据的人来说，理解傅里叶变换的底层逻辑至关重要。它能帮你更好地解读频谱图，选择合适的窗函数，甚至诊断算法中的诡异问题。我们将从最直观的DFT定义式出发，编写代码，然后与NumPy的FFT进行性能对决，最后剖析FFT算法如何将计算复杂度从O(N²)降到O(N log N)。准备好了吗？让我们开始吧。

1. 离散傅里叶变换：从数学公式到Python代码

傅里叶变换的核心思想，是将一个随时间变化的信号，分解成一系列不同频率、不同振幅的正弦波之和。在连续世界中，这涉及积分运算。但计算机处理的是离散的采样点，因此我们需要的是离散傅里叶变换。

DFT的公式看起来并不复杂。对于一个长度为N的复数序列x[n]，其离散傅里叶变换X[k]定义为：

X[k] = Σ_{n=0}^{N-1} x[n] * exp(-i * 2π * k * n / N)

其中，k = 0, 1, ..., N-1 代表频率索引，i是虚数单位。这个公式的直观解释是：对于每一个目标频率k，我们将原始信号x[n]与一个该频率的复正弦波（exp(-i * 2π * k * n / N)）相乘，并对所有时间点n求和。结果X[k]是一个复数，其模长代表该频率成分的振幅，辐角代表该频率成分的相位。

注意：公式中的指数项 exp(-i * 2π * k * n / N) 就是欧拉公式的体现，它等价于 cos(2π * k * n / N) - i * sin(2π * k * n / N)。这构成了一个在复平面上旋转的单位向量。

根据这个定义，我们可以直接翻译成Python代码。这就是所谓的“朴素DFT”实现，它完全遵循数学定义，没有任何优化。

import numpy as np
import time

def naive_dft(x):
    """
    朴素离散傅里叶变换实现。
    参数:
        x: 一维numpy数组，输入信号（可以是实数或复数）。
    返回:
        X: 一维复数numpy数组，DFT结果。
    """
    N = len(x)
    # 初始化一个全零的复数数组来存放结果
    X = np.zeros(N, dtype=np.complex128)
    
    # 遍历每个输出频率点 k
    for k in range(N):
        # 计算旋转因子（twiddle factors）的指数部分
        # 这里我们预先计算 (2π * k / N)，避免在内部循环重复计算
        base = -2j * np.pi * k / N
        # 对每个时间点 n 进行求和
        for n in range(N):
            X[k] += x[n] * np.exp(base * n)
    return X

这段代码非常直观，但它暴露了一个严重的问题：双重嵌套的for循环。这意味着计算量随着数据点N的增加呈平方级增长。我们来分析一下它的计算复杂度：

• 乘法次数：对于每个k（外循环N次），对于每个n（内循环N次），我们执行一次复数乘法和一次复数加法。所以总共是大约 N * N = N² 次复数运算。
• 指数计算：np.exp(base * n) 在每次内循环都会被计算，这是一个相对昂贵的运算。尽管我们做了一点优化，将base提到外循环，但指数计算本身依然很耗时。

为了验证我们的实现是否正确，我们可以用一个简单的信号来测试：一个由两个正弦波叠加而成的信号。

def test_naive_dft():
    # 生成一个测试信号：1Hz正弦波 + 3Hz正弦波
    N = 32  # 采样点数
    t = np.linspace(0, 1, N, endpoint=False)  # 1秒时长
    freq1, freq2 = 1, 3
    signal = np.sin(2 * np.pi * freq1 * t) + 0.5 * np.sin(2 * np.pi * freq2 * t)
    
    # 使用我们的朴素DFT
    my_fft = naive_dft(signal)
    
    # 使用NumPy的FFT作为基准
    np_fft = np.fft.fft(signal)
    
    # 比较结果（允许微小的数值误差）
    print("最大绝对误差:", np.max(np.abs(my_fft - np_fft)))
    # 通常误差在1e-12到1e-14量级，证明算法正确
    
    # 绘制幅度谱对比
    import matplotlib.pyplot as plt
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    ax1.stem(np.abs(my_fft[:N//2]))
    ax1.set_title('自实现DFT幅度谱')
    ax1.set_xlabel('频率索引 k')
    ax1.set_ylabel('振幅')
    
    ax2.stem(np.abs(np_fft[:N//2]))
    ax2.set_title('NumPy FFT幅度谱')
    ax2.set_xlabel('频率索引 k')
    ax2.set_ylabel('振幅')
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    test_naive_dft()

运行这段代码，你会看到两个频谱图几乎完全一致，误差小到可以忽略不计。这证明我们的朴素DFT在数学上是正确的。然而，当我们尝试增加数据量时，问题就来了。

2. 性能对决：朴素DFT vs NumPy FFT

理论上的复杂度分析需要实际的性能测试来验证。让我们设计一个实验，测量两种算法在不同数据规模下的耗时。

def performance_comparison():
    sizes = [16, 32, 64, 128, 256, 512]  # 测试的数据规模
    naive_times = []
    numpy_times = []
    
    for N in sizes:
        # 生成随机信号
        x = np.random.randn(N) + 1j * np.random.randn(N)
        
        # 测试朴素DFT
        start = time.perf_counter()
        _ = naive_dft(x)
        naive_time = time.perf_counter() - start
        naive_times.append(naive_time)
        
        # 测试NumPy FFT
        start = time.perf_counter()
        _ = np.fft.fft(x)
        numpy_time = time.perf_counter() - start
        numpy_times.append(numpy_time)
        
        print(f"N={N:3d} | 朴素DFT: {naive_time:.6f}s | NumPy FFT: {numpy_time:.6f}s | 加速比: {naive_time/numpy_time:.1f}x")
    
    # 绘制性能对比图
    import matplotlib.pyplot as plt
    plt.figure(figsize=(10, 6))
    plt.plot(sizes, naive_times, 'o-', label='朴素DFT (O(N²))', linewidth=2)
    plt.plot(sizes, numpy_times, 's-', label='NumPy FFT (O(N log N))', linewidth=2)
    plt.xlabel('数据规模 N')
    plt.ylabel('计算时间 (秒)')
    plt.title('DFT算法性能对比 (对数坐标)')
    plt.yscale('log')
    plt.xscale('log')
    plt.grid(True, which="both", ls="--", alpha=0.5)
    plt.legend()
    plt.show()

performance_comparison()

运行这个对比测试，你会看到类似下面的结果（具体时间取决于你的硬件）：

这个表格清晰地揭示了一个事实：随着N的增大，朴素DFT的时间增长极快，而NumPy FFT的时间增长则缓慢得多。当N=512时，FFT已经比朴素实现快了上万倍。这种差距正是算法优化的威力所在。

提示：在实际项目中，数据规模动辄成千上万（音频处理）甚至数百万（图像处理）。如果使用O(N²)的算法，计算将变得完全不可行。这就是为什么FFT算法被认为是20世纪最重要的算法之一。

那么，FFT究竟施展了什么魔法？它如何将计算量从N²降到N log N？接下来，我们就揭开这个魔术的幕布。

3. FFT的核心思想：分而治之与旋转因子

FFT算法有很多变种，最著名的是Cooley-Tukey算法。它的核心思想是分治法——将一个大的DFT分解成多个小的DFT。这里我们以最常见的基2-FFT（要求N是2的整数次幂）为例，看看它是如何工作的。

首先，我们回顾一下DFT公式：

X[k] = Σ_{n=0}^{N-1} x[n] * W_N^{kn}

其中 W_N = exp(-i * 2π / N) 被称为旋转因子。注意W_N^{kn}具有周期性：W_N^{k(n+N)} = W_N^{kn}，W_N^{(k+N)n} = W_N^{kn}。

FFT的巧妙之处在于，当N是偶数时，我们可以将输入序列x[n]按奇偶索引拆分成两个子序列：

• 偶数索引序列：x_e[m] = x[2m], 其中 m = 0, 1, ..., N/2-1
• 奇数索引序列：x_o[m] = x[2m+1], 其中 m = 0, 1, ..., N/2-1

然后，原DFT可以重新表达为这两个长度为N/2的序列的DFT的组合：

X[k] = Σ_{m=0}^{N/2-1} x[2m] * W_N^{k*(2m)} + Σ_{m=0}^{N/2-1} x[2m+1] * W_N^{k*(2m+1)}
     = Σ_{m=0}^{N/2-1} x_e[m] * (W_N^2)^{km} + W_N^k * Σ_{m=0}^{N/2-1} x_o[m] * (W_N^2)^{km}

注意到 W_N^2 = exp(-i * 2π * 2 / N) = exp(-i * 2π / (N/2)) = W_{N/2}。这正是长度为N/2的DFT的旋转因子！

因此，我们得到关键的递归关系：

X[k] = X_e[k] + W_N^k * X_o[k]

其中X_e[k]和X_o[k]分别是偶数子序列和奇数子序列的DFT结果（长度为N/2）。而且，由于DFT的周期性，对于k >= N/2的部分，我们有：

X[k + N/2] = X_e[k] - W_N^k * X_o[k]

看，一个长度为N的DFT，被分解成了两个长度为N/2的DFT，再加上一些乘法和加法的组合操作。如果N/2仍然是偶数，我们可以继续分解下去，直到子问题的长度为1（长度为1的序列的DFT就是它本身）。这就是分治法的精髓。

让我们用代码实现这个递归版本的FFT：

def recursive_fft(x):
    """
    递归实现的Cooley-Tukey基2-FFT算法。
    要求输入长度N是2的整数次幂。
    """
    N = len(x)
    
    # 基线条件：长度为1的DFT就是它本身
    if N == 1:
        return x
    
    # 检查N是否为2的幂
    if N & (N - 1) != 0:
        raise ValueError("输入长度必须是2的整数次幂")
    
    # 将输入按奇偶索引拆分
    x_even = recursive_fft(x[0::2])  # 偶数索引
    x_odd = recursive_fft(x[1::2])   # 奇数索引
    
    # 初始化结果数组
    X = np.zeros(N, dtype=np.complex128)
    
    # 计算旋转因子
    # 我们只需要计算前N/2个，因为后半部分可以利用对称性
    W_N = np.exp(-2j * np.pi / N * np.arange(N//2))
    
    # 组合结果
    for k in range(N//2):
        t = W_N[k] * x_odd[k]
        X[k] = x_even[k] + t
        X[k + N//2] = x_even[k] - t
    
    return X

这个递归实现已经比朴素DFT快了很多，但递归调用本身有开销。在实际的库实现中（如NumPy），使用的是迭代版本的FFT，它通过位反转置换和蝴蝶操作来实现，避免了递归调用栈的开销，并且能更好地利用CPU缓存。

为了帮助理解从递归到迭代的转换，我们来看一下N=8时的计算流程图（蝴蝶图）：

阶段1 (长度2 DFT)       阶段2 (长度4 DFT)       阶段3 (长度8 DFT)
x[0]  ────────────┬─────────── X[0]
x[4]  ──────┐     │     ┌───── X[1]
x[2]  ───┐  │  ┌──┴──┐  │  ┌── X[2]
x[6]  ─┐ │  │  │  │  │  │  │  ─ X[3]
x[1]  ─┼─┼──┼──┼──┼──┼──┼──┼── X[4]
x[5]  ─┘ │  │  │  │  │  │  │  ─ X[5]
x[3]  ───┘  │  │  │  │  │  │  ─ X[6]
x[7]  ──────┘  │  │  │  │  │  ─ X[7]

图中的每一条连线都代表一次复数乘加运算（蝴蝶操作），每个阶段有N次操作，总共有log₂N个阶段，所以总操作次数约为 N * log₂N。与N²相比，当N很大时，节省的计算量是惊人的。

4. 工程实践：优化实现与逆变换

理解了FFT的原理后，我们可以尝试实现一个更高效的迭代版本。这个版本会涉及到位反转——这是将递归计算转换为迭代计算的关键步骤。

def iterative_fft(x):
    """
    迭代实现的基2-FFT算法（非原位计算，易于理解）。
    输入x可以是实数或复数，长度必须是2的幂。
    """
    N = len(x)
    
    # 位反转置换
    # 将输入重新排列成递归FFT最底层的顺序
    j = 0
    for i in range(1, N):
        bit = N >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j ^= bit
        if i < j:
            x[i], x[j] = x[j], x[i]
    
    # 迭代进行蝴蝶操作
    length = 2
    while length <= N:
        half_len = length // 2
        # 计算本层使用的旋转因子
        W_length = np.exp(-2j * np.pi / length)
        
        for start in range(0, N, length):
            w = 1.0 + 0.0j
            for k in range(half_len):
                # 蝴蝶操作
                u = x[start + k]
                v = x[start + k + half_len] * w
                x[start + k] = u + v
                x[start + k + half_len] = u - v
                # 更新旋转因子
                w *= W_length
        length <<= 1  # 长度翻倍
    
    return x

def my_fft(x):
    """对iterative_fft的包装，处理非2的幂长度（通过补零）"""
    N_orig = len(x)
    # 找到下一个2的幂
    N_pow2 = 1
    while N_pow2 < N_orig:
        N_pow2 <<= 1
    
    # 如果长度不是2的幂，补零
    if N_pow2 != N_orig:
        x_padded = np.zeros(N_pow2, dtype=np.complex128)
        x_padded[:N_orig] = x
        result = iterative_fft(x_padded)
        return result[:N_orig]  # 通常只返回前N_orig个点
    else:
        return iterative_fft(x.copy())  # 复制以避免修改原数组

现在，我们有了自己的FFT实现。但傅里叶变换通常是成对使用的：正变换将信号从时域转到频域，逆变换则从频域恢复时域信号。逆FFT与正FFT非常相似，只是旋转因子的指数符号相反，并且结果需要除以N。

def my_ifft(X):
    """
    基于FFT的逆FFT实现。
    利用性质：IFFT(x) = (1/N) * conj(FFT(conj(x)))
    其中conj表示取共轭。
    """
    N = len(X)
    # 计算共轭
    X_conj = np.conj(X)
    # 对共轭进行FFT
    x_conj_fft = my_fft(X_conj)
    # 再次取共轭并除以N
    x = np.conj(x_conj_fft) / N
    return x

# 测试正逆变换的准确性
def test_fft_ifft_roundtrip():
    # 生成一个复杂的测试信号
    N = 64
    t = np.linspace(0, 1, N, endpoint=False)
    # 包含多个频率成分
    signal = (np.sin(2 * np.pi * 5 * t) + 
              0.5 * np.cos(2 * np.pi * 12 * t) +
              0.3 * np.sin(2 * np.pi * 20 * t + np.pi/4))
    
    # 使用我们的实现
    my_spectrum = my_fft(signal)
    my_reconstructed = my_ifft(my_spectrum)
    
    # 使用NumPy作为基准
    np_spectrum = np.fft.fft(signal)
    np_reconstructed = np.fft.ifft(np_spectrum)
    
    # 比较重建误差
    error_my = np.max(np.abs(signal - my_reconstructed))
    error_np = np.max(np.abs(signal - np_reconstructed))
    
    print(f"自实现FFT/IFFT往返最大误差: {error_my:.2e}")
    print(f"NumPy FFT/IFFT往返最大误差: {error_np:.2e}")
    
    # 可视化
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    
    axes[0, 0].plot(t, signal, 'b-', label='原始信号')
    axes[0, 0].set_title('原始时域信号')
    axes[0, 0].set_xlabel('时间')
    axes[0, 0].set_ylabel('振幅')
    axes[0, 0].legend()
    axes[0, 0].grid(True)
    
    axes[0, 1].plot(t, my_reconstructed.real, 'r--', label='重建信号', alpha=0.7)
    axes[0, 1].plot(t, signal, 'b-', label='原始信号', alpha=0.3)
    axes[0, 1].set_title('自实现重建信号对比')
    axes[0, 1].set_xlabel('时间')
    axes[0, 1].set_ylabel('振幅')
    axes[0, 1].legend()
    axes[0, 1].grid(True)
    
    # 幅度谱对比
    freq = np.fft.fftfreq(N, d=t[1]-t[0])
    axes[1, 0].stem(freq[:N//2], np.abs(my_spectrum[:N//2]), 
                    linefmt='b-', markerfmt='bo', basefmt=' ', label='自实现')
    axes[1, 0].set_title('自实现FFT幅度谱')
    axes[1, 0].set_xlabel('频率 (Hz)')
    axes[1, 0].set_ylabel('振幅')
    axes[1, 0].grid(True)
    
    axes[1, 1].stem(freq[:N//2], np.abs(np_spectrum[:N//2]), 
                    linefmt='r-', markerfmt='ro', basefmt=' ', label='NumPy', alpha=0.7)
    axes[1, 1].set_title('NumPy FFT幅度谱')
    axes[1, 1].set_xlabel('频率 (Hz)')
    axes[1, 1].set_ylabel('振幅')
    axes[1, 1].grid(True)
    
    plt.tight_layout()
    plt.show()

test_fft_ifft_roundtrip()

这个完整的测试展示了从正变换到逆变换的完整流程。你会看到，即使是我们自己实现的FFT，重建误差也非常小（通常在10^-14量级），这证明了算法的正确性。而NumPy的实现由于使用了更底层的优化（可能是用C或Fortran写的，并且可能使用了SIMD指令），误差通常更小，速度也更快。

在实际项目中，理解这些底层原理的价值在于，当遇到奇怪的现象时（比如频谱泄漏、栅栏效应），你知道问题可能出在哪里。例如，如果你发现频谱不对称，可能是没有正确理解实信号DFT的共轭对称性；如果重建信号有偏差，可能是缩放因子（1/N）放错了位置。

手写DFT和FFT的过程，就像亲手搭建一座房子，而不是直接住进精装公寓。你知道每一根梁柱在哪里，知道电路和水管如何排布。当房子出现问题时，你清楚地知道该检查哪里。在信号处理的世界里，这种深度的理解是无价的。下次当你调用np.fft.fft时，你看到的将不再是一个黑盒魔法，而是一个优雅的分治算法在高效运转。