python实现快速傅里叶变换（FFT）

 1、利用矩阵乘（速度慢）

import numpy as np

"""
计算离散傅里叶变换
"""
def dft(x):
    x = np.asarray(x, dtype=float)                                  # 将输入的x转换为浮点数
    N = x.shape[0]                                                  # 输出x的形状(1024,)中的第一个位置1024
    n = np.arange(N)                                                # 输出数组[0, 1, 2, ..., 1023]
    k = n.reshape((N, 1))                                           # 将1行1024列转换为1024行1列
    M = np.exp(-2j * np.pi * k * n / N)                             # 二维数组（复数形式表示）（相当于坐标轴）

    return np.dot(M, x)                                             # 返回二维数组与x的点成

"""
递归计算快速傅里叶变换（奇偶分离）
"""
def fft(x):
    x = np.asarray(x, dtype=float)
    N = x.shape[0]
    if N % 2 > 0:                                                   # N不是2的幂
        raise ValueError("must be an power of 2")
    elif N <= 2 :                                                   # N小于等于2
        return dft(x)
    else:                                                           # N大于2并且为2的倍数
        X_even = fft(x[::2])                                        # 递归得到下标为偶数的fft，总数为N/2个
        X_odd = fft(x[1::2])                                        #              奇数            N/2个
        terms = np.exp(-2j * np.pi * np.arange(N) /N)               # 基矩阵M，M里面的值有N个
        return np.concatenate([X_even + terms[:int(N/2)] * X_odd,   # 返回前N/2个
                               X_even + terms[int(N/2):] * X_odd])  # 返回后N/2个

x = np.random.random(8)
print(fft(x))
print(np.fft.fft(x))
print(np.allclose(fft(x), np.fft.fft(x)))                           # fft与numpy.fft分别计算，值是否对应相等，默认在1e-05的误差范围内

运行结果：

[ 3.22533499+0.00000000e+00j  0.80576208+5.28570440e-01j
 -0.957105  +5.98902902e-02j -0.34310532+7.70665268e-02j
  0.15598139-1.87943703e-16j -0.34310532-7.70665268e-02j
 -0.957105  -5.98902902e-02j  0.80576208-5.28570440e-01j]
[ 3.22533499+0.j          0.80576208+0.52857044j -0.957105  +0.05989029j
 -0.34310532+0.07706653j  0.15598139+0.j         -0.34310532-0.07706653j
 -0.957105  -0.05989029j  0.80576208-0.52857044j]
True

Process finished with exit code 0

%timeit dft(x)
%timeit fft(x)
%timeit np.fft.fft(x)

 

2、利用向量乘（速度快）

def fft_v(x):
    x = np.asarray(x, dtype=float)
    N = x.shape[0]
    if np.log2(N) % 1 > 0:
        raise ValueError("must be a power of 2")
        
    N_min = min(N, 2)
    
    n = np.arange(N_min)
    k = n[:, None]
    M = np.exp(-2j * np.pi * n * k / N_min)
    X = np.dot(M, x.reshape((N_min, -1)))
while X.shape[0] < N:
        X_even = X[:, :int(X.shape[1] / 2)]
        X_odd = X[:, int(X.shape[1] / 2):]
        terms = np.exp(-1j * np.pi * np.arange(X.shape[0])
                        / X.shape[0])[:, None]
        X = np.vstack([X_even + terms * X_odd,
                       X_even - terms * X_odd])
return X.ravel()

x = np.random.random(1024)
np.allclose(fft_v(x), np.fft.fft(x))

True

%timeit fft(x)
%timeit fft_v(x)
%timeit np.fft.fft(x)