python实现DIT-FFT

1、导入相应的包

import math
import numpy as np 
from numpy import cos,sin,pi
from scipy.fftpack import fft,ifft

2、 定义 FFT

def my_fft(arr):
    length = len(arr)
    # length == int(2**math.log(length,2)) 是错误的判断
    # M 为级数
    M = math.log(length, 2) if length == 2**math.log(
        length, 2) else int(math.log(length, 2)) + 1
    # 注意要将 M 转为 int 类型
    M = int(M)
    # 填充序列
    for i in range(len(arr), int(2**M)):
        arr.append(0)
    # 更新序列的长度
    N = len(arr)
    # 中间数组，
    array = [i for i in range(N)]

    def reverse(a):  # 二进制倒序算法
        L = M
        for i in range(N):
            a[i] = (a[i >> 1] >> 1) | ((i & 1) << (L - 1))

    def butterfly(xn, x1_point, x2_point, Wn):  #  蝶形计算单元

        T = xn[x2_point] * Wn
        # TMD 搞心态 result2 = xn[x2_point] - T
        result1 = xn[x1_point] + T
        result2 = xn[x1_point] - T
        xn[x1_point] = result1
        xn[x2_point] = result2

    def creat_wn(m, j):
        m = 2**m
        return cos(2 * pi * m * j / N) - sin(2 * pi * m * j / N) * 1j

    def single__fft(m, xn):  # 单级fft

        point_distance = pow(2, m - 1)  # 点距离 形如 x(0) - x(4)
        group_distance = 2 * point_distance  # 组别距离
        group_count = length // group_distance  # 组别数量
        group_header = 0
        for i in range(group_count):
            group_header = i * group_distance  #每一组的第一个元素
            for j in range(group_distance // 2):
                # 特别要注意的是每一级Wn 的下标是反着来的 即
                '''
                
                级数    1             2             3 ....             M
                下标    2^(M)      2^(M-1)          2^(M-2)            2^0
                
                可得第i级 Wx 的下标为 x = N/(2**(M-i))
                '''
#               旋转因子的上标可知为 组别距离除以2 或者直接是点距离
#               另外最后一次的组别距离为N,即所有的输入为一个组
                w = creat_wn(M - m, j)
#               print(j, "****", w)
                x1 = j + group_header
                x2 = x1 + point_distance
                butterfly(xn, x1, x2, w)
#               中间步骤可以通过打印
#               print(xn)

    reverse(array)
    new_array = []
    for i in range(length):
        new_array.append(arr[array[i]])
    for i in range(1, M + 1):
        single__fft(i, new_array)
    print(new_array)

3、测试

arr = [1,3,2,1]
my_fft(arr)
# 调用包中的fft对比结果
fft(arr)

4、 测试结果

本文思路部分参考