基於Python實現DIT-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(jomega)|$')
    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(jomega)|$')
    plt.grid(linestyle="-.")
    plt.legend()
    plt.show()

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(jomega)|$')
    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