自己写函数实现fft
使用递归方法
from math import log, ceil, cos, sin, pi
import matplotlib.pyplot as plt
import numpy as np
# 这两行代码解决 plt 中文显示的问题
plt.rcparams['font.sans-serif'] = ['simhei']
plt.rcparams['axes.unicode_minus'] = false
def fft(x, n=none):
# dit-fft 函数说明
# x: 时域序列
# n: n点dft, 理论上n=2**m
# 返回值为序列x的dft
if n is none:
n = len(x)
elif n < len(x):
n = len(x)
if n == 2:
return [x[0]+x[1], x[0]-x[1]]
# 补0使得n=2**m
m = ceil(log(n, 2))
n = 2**m
x = x + [0] * (n-len(x))
# 递归地计算偶数项和奇数项的dft
x1 = fft(x[0::2])
x2 = fft(x[1::2])
x = [0] * n
for i in range(n//2):
# 蝶形计算
tmp = (cos(2*pi/n*i)-1j*sin(2*pi/n*i))*x2[i]
x[i] = x1[i] + tmp
x[i+n//2] = x1[i] - tmp
return x
if __name__ == '__main__':
x = [1]*10
y = fft(x, 1024)
# print(y)
z = [abs(i) for i in y]
# print(z)
plt.plot(np.arange(len(z))*2/len(z), z, label='10点矩形窗函数的fft')
plt.title("幅度谱")
plt.xlabel(r'单位：$\pi$')
plt.ylabel(r'$|h(j\omega)|$')
plt.grid(linestyle="-.")
plt.legend()
plt.show()
使用循环，流式计算（极大地节省了内存）
from math import log, ceil, cos, sin, pi
import matplotlib.pyplot as plt
import numpy as np
# 这两行代码解决 plt 中文显示的问题
plt.rcparams['font.sans-serif'] = ['simhei']
plt.rcparams['axes.unicode_minus'] = false
def fft(x, n=none):
# dit-fft 函数说明
# x: 时域序列
# n: n点dft, 理论上n=2**m
# 返回值为序列x的dft
"""
采用流式计算方法，只用了一个n(n=2**m)点的数组内存
"""
if n is none:
n = len(x)
elif n < len(x):
n = len(x)
# 补0使得：n=2**m
m = ceil(log(n, 2))
n = 2**m
x = x + [0] * (n-len(x))
fm = "{:0"+f"{m}"+"b}"
x = [0] * n
for i in range(n//2):
index1 = eval('0b'+fm.format(i*2)[::-1])
index2 = eval('0b'+fm.format(i*2+1)[::-1])
x[2*i] = x[index1] + x[index2]
x[2*i+1] = x[index1] - x[index2]
for i in range(1, m):
# 第i步表示将2**i点dft合成2**(i+1)点的dft
# 蝶形宽度width
width = 2**i
"""
将x(k)序列进行分组，每组2**(i+1)个点，
便于将每组中两组2**i点dft合成一组2**(i+1)点的dft
"""
# num=2*width=2**(i+1), 表示每组点数
num = 2*width
# 组数groups
groups = n//num
for j in range(groups):
# 对每组将2**i点dft合成2**(i+1)=num点的dft
for k in range(num//2):
# 旋转因子
w = cos(2*pi/num*k) - 1j * sin(2*pi/num*k)
# 第j组第k个
index = j*num + k
tmp = w * x[index+width] # 每个蝶形一次复数乘法
x[index], x[index+width] = x[index]+tmp, x[index]-tmp
return x
if __name__ == '__main__':
x = [1]*10
y = fft(x, 1024)
# print(y)
z = [abs(i) for i in y]
# print(z)
plt.plot(np.arange(len(z))*2/len(z), z, label='10点矩形窗函数的fft')
plt.title("幅度谱")
plt.xlabel(r'单位：$\pi$')
plt.ylabel(r'$|h(j\omega)|$')
plt.grid(linestyle="-.")
plt.legend()
plt.show()
运行结果：
# 说明：建议使用第二种方法实现fft。第一种递归的方法在递归调用时也需要一定的成本，且使用的内存较大；而第二种方法只使用了一个n(n=2**m)点的数组进行计算，内存可重用。
使用python的第三方库进行fft
import numpy as np
from numpy.fft import fft, ifft
# from scipy.fftpack import fft, ifft
import matplotlib.pyplot as plt
# 这两行代码解决 plt 中文显示的问题
plt.rcparams['font.sans-serif'] = ['simhei']
plt.rcparams['axes.unicode_minus'] = false
if __name__ == '__main__':
x = 2*np.sin(np.pi/2*np.arange(100))+np.sin(np.pi/5*np.arange(100))
z = [abs(i) for i in fft(x, 2048)]
# print(z)
l = len(z)
plt.plot((np.arange(l)*2/l)[:l//2], z[:l//2], label='两个不同频率正弦信号相加的dft')
plt.title("幅度谱")
plt.xlabel('$\pi$')
plt.ylabel('$|h(j\omega)|$')
plt.grid(linestyle="-.")
plt.legend()
plt.show()
print('max(abs(ifft(fft(x))-x)) = ', end='')
print(max(abs(ifft(fft(x))-x)))
运行结果：
max(abs(ifft(fft(x))-x)) = 9.01467522361575e-16
以上就是基于python实现dit-fft算法的详细内容，更多关于python dit-fft算法的资料请关注其它相关文章！
