基于Python实现DIT-FFT算法

Tertia ·

更新时间:2024-09-21

· 1378 次阅读

自己写函数实现FFT

使用python的第三方库进行FFT

自己写函数实现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)))

运行结果：