Python实现离散傅里叶变换（DFT）出现无限循环问题求助

问题说明 #

编写Python实现离散傅里叶变换（DFT）时，代码看似陷入无限循环，必须键盘中断才能停止，但实际是运算量过大导致的假死。以下是你的代码及解决思路：

原代码 #

import matplotlib.pyplot as plt
import numpy as np
import cmath

Fs = 40000;  # 采样率
Ts = 1.0/Fs; # 采样周期
t = np.arange(0,1,Ts) # 时间向量

f = 100;   # 信号频率
x1_n = np.sin(2*np.pi*f*t + 0)
f = 1000;
x2_n = np.sin(2*np.pi*f*t + 180)

x_n = x1_n + x2_n

n = len(x_n) # 信号长度
k = np.arange(n) # k向量
T = n/Fs
frq = k/T # 双边频率向量
frq = frq[range(int(n/2))] # 单边频率向量

X = np.fft.fft(x_n)/n # numpy的FFT结果（用于验证）

print("A")
print(X) # 打印numpy FFT结果

m = len(x_n)
output = []
for k in range(m):  # 遍历每个输出点
    s = complex(0)
    for t in range(m):  # 遍历每个输入点
        angle = 2j * cmath.pi * t * k / m
        s += x_n[t] * cmath.exp(-angle)
    output.append(s)

print("B") # 打印手动实现的DFT结果
print(output)

核心问题 #

你的代码没有无限循环，但信号长度m = len(x_n) = 40000，双重循环需要执行40000×40000=16亿次运算，普通电脑需要很长时间才能完成，所以看起来像卡住了。

解决思路 #

1. 先缩小测试规模验证逻辑 #

先减小信号长度，快速验证代码逻辑是否正确：

• 修改时间向量为t = np.arange(0, 0.01, Ts)，这样信号长度m=400，运算量降到16万次，几秒就能出结果。
• 验证时记得给手动实现的结果做归一化（除以m），才能和numpy的X = np.fft.fft(x_n)/n结果对齐：

  # 在output.append(s)之后添加
  output = [s/m for s in output]
  

2. 优化代码效率 #

如果需要处理大规模信号，纯Python双重循环的效率太低，推荐两种优化方式：

• numpy矢量化实现：利用numpy的广播机制替代循环，速度比纯Python循环快几十倍：

  m = len(x_n)
  k = np.arange(m)
  t = np.arange(m)
  # 生成角度矩阵（利用广播）
  angle = 2j * np.pi * np.outer(t, k) / m
  # 矢量化计算DFT
  output = np.sum(x_n[:, np.newaxis] * np.exp(-angle), axis=0)
  output = output / m  # 归一化
  

• 实现FFT算法：DFT的时间复杂度是O(n²)，而FFT是O(n log n)，处理大信号时效率天差地别，比如实现基2FFT算法（要求信号长度是2的幂次）。

3. 细节修正 #

• 代码中同时使用了cmath.pi和np.pi，二者数值一致，建议统一用其中一个，避免混淆。
• 注意相位的单位：x2_n = np.sin(2*np.pi*f*t + 180)中的180是角度，而numpy的三角函数用弧度，应该改成np.pi（180°=π弧度），否则信号相位不符合预期。

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