仿照 c 语言,有很多丑陋的 for 循环
def __fft(sig, inv):
N = sig.size # 原始长度,如果不是二的幂,补零
bit=0 # 表示信号长度的最小比特数
while ((1<<bit)<N):
bit += 1;
N = pow(2,bit) # 不小于 N 的最小二次幂
sig = np.hstack((sig,np.zeros((sig.shape[0],N-sig.shape[1]),dtype=complex))) # 补零
# 重新排序
rev = [0 for i in range(N)]
for i in range(N):
rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (bit - 1))
if(i<rev[i]):
temp = sig[0,i]
sig[0,i] = sig[0,rev[i]]
sig[0,rev[i]] = temp
# 蝶形运算
freq = sig.astype(complex)
for mid in [pow(2,i) for i in range(bit)]:
for i in range(0,N,2*mid):
for j in range(mid):
x = freq[0,i+j]
y = freq[0,i+j+mid]*np.exp(-1j*inv*np.pi*(i+j)/mid)
freq[0,i+j] = x+y
freq[0,i+j+mid] = x-y
if inv == -1:
return freq/N
return freq
测试
n = 16
sig = np.matrix([ np.random.random() for i in range(pow(2,n))]) # 随机信号
cov2 = fft(sig) # numpy.fft 得到的结果
freq = __fft(sig,1)
print('distance from numpy.fft: ',np.linalg.norm(freq[0,:cov2.size]-cov2))
back = __fft(freq,-1) # 逆变换
print('reconstruction error: ',np.linalg.norm(sig-back[0,:cov2.size]))
结果
distance from numpy.fft: 2.5842860851574793e-07
reconstruction error: 1.0518201158802643e-09
上面第一版实在是太丑陋了,因为 python 根本不适合跑 for 循环
而且,在 python 数组中,可以利用切片非常方便地取出 奇数列 和 偶数列
用递归简化代码
def dft(sig):
N = sig.size
V = np.matrix([[np.exp(-1j*2*np.pi*v*y/N) for v in range(N)] for y in range(N)])
return sig.dot(V)
def FFT(x):
"""A recursive implementation of the 1D Cooley-Tukey FFT"""
x = np.asarray(x, dtype=float)
N = x.shape[0]
if N % 2 > 0:
raise ValueError("size of x must be a power of 2")
elif N <= 8: # this cutoff should be optimized
return dft(x)
else:
X_even = FFT(x[::2])
X_odd = FFT(x[1::2])
# print('size X_even',X_even.shape)
factor = np.matrix(np.exp(-2j * np.pi * np.arange(N) / N))
# print('size factor',factor.shape)
return np.hstack([X_even + np.multiply(factor[0,:int(N/2)],X_odd),
X_even + np.multiply(factor[0,int(N/2):],X_odd)])
测试
n = 16
sig = [ np.random.random() for i in range(pow(2,n))]
cov2 = fft(sig)
freq = FFT(sig)
print('distance from numpy.fft: ',np.linalg.norm(freq-cov2))
distance from numpy.fft: 3.946983495917491e-11