Python DFT、功率谱估计

x(n)=10sin(2pi0.1+3/pi)+4sin(2pi0.4+4/pi),噪声为均值为零方差为一的高斯随机噪声,DFT点数N=1024,目标为画出频谱和功率谱，并对f1和f2（0.1和0.4）进行估计。在估计时，将对功率谱拟合和对DFT拟合比较后选择了对DFT拟合

最终结果为：

对功率谱拟合后得到的结果
对DFT拟合后得到的结果

代码如下

import matplotlib.pyplot as plt
import numpy as np
from math import *
import random as ran
#设定初始参数
f1 = 0.1
f2 = 0.4
N = 1024 
#生成空列表存储结果
Xk = []
power = [] 
#定义Wn
def WN(N,k,n):
    return complex(cos(-2*pi/N*k*n), sin(-2*pi/N*k*n)) 
#计算DFT和功率谱
sum = 0
for k in range(0,N):
    for n in range(0,N):
      xn = 10*sin(2*pi*f1*n+3/pi)+4*sin(2*pi*f2*n+4/pi)+ran.gauss(0,1)
      sum += xn*WN(N, k, n)
    Xk.append(sum)
    power.append(sum**2)
    sum = 0 
#画图：DFT，功率谱（实际上为散点图，画为连续曲线）
plt.subplot(1, 3, 1)
plt.title('DFT')
plt.plot(np.abs(Xk))
plt.subplot(1, 3, 2)
plt.title('power spectrum')
plt.plot(np.abs(power)/N)
#取DFT前半部分结果进行拟合并画图（画图是为了方便观察拟合结果并进行调整）
for i in range(0, int(N/2)):
    Xk.pop()
Xk = list(np.abs(Xk))
plt.subplot(1, 3, 3)
x = [i for i in range(0, int(N/2))]
y = np.polyfit(x,Xk,6)
Xk_smooth = np.polyval(y,x)
plt.title('Xk_smooth')
plt.plot(Xk_smooth)
plt.show()
#对f1、f2进行估计得到对应的k值并存入列表F（方便最终输出）：波峰处值最大=对于任意的k，X(k)比两侧大即为波峰
F=[]
for l in range(1, int(N/2)-1):
    if Xk_smooth[l] > Xk_smooth[l-1] and Xk_smooth[l] > Xk_smooth[l+1]:
        F.append(l)
#对F进行处理得到估计的f1、f2并输出
t = 0
for f in F:
    f = f/N
    t += 1
    print('f'+str(t)+'='+str(f))