快速傅里叶变换(FFT)是信号处理和数据分析中最重要的算法之一。我虽然已经使用了多年,但是没有正式的计算机科学背景,本周我发现我从未想过FFT 如何快速地计算离散傅立叶变换。我翻开尘封已久的算法书开始研究JW Cooley和John Tukey在其1965年的论文中提出的这个简单但令人困惑的计算技巧。
这篇文章的目标是深入研究Cooley-Tukey FFT算法,解释它如何利用对称性完成DFT的优化,并使用Python实现算法,将理论付诸实践。我希望本文能让像我这样的数据科学家更全面地了解我们使用的算法背后发生的事情。
FFT是一个快速的算法用于计算离散傅里叶变换(DFT)。该算法将离散傅里叶变换
正向离散傅里叶变换(DFT)
反向离散傅里叶变换(IDFT)
从
由于FFT在很多领域中十分重要,Python有许多标准工具和包用于计算FFT。NumPy和SciPy都有通过完整测试的FFTPACK库的包装,可以在numpy.fft
和scipy.fftpack
子模块中找到。据我所知,最快的FFT是FFTW包。可以通过安装PyFFTW包在Python中使用它。
现在,我们先把这些实现放一边,来研究如何从头开始在Python中实现FFT。
为了简单起见,我们只考虑正向变换,因为反向变换可以以非常相似的手段实现。观察上面DFT的公式,我们可以看到这只不过是一个直接的线性变换操作:一个矩阵-向量乘法
其中矩阵M由以下计算得
使用上式思想,我们就可以使用下面简单的矩阵乘法计算DFT:
import numpy as np
def DFT_slow(x):
"""Compute the discrete Fourier Transform of the 1D array x"""
x = np.asarray(x, dtype=float)
N = x.shape[0]
n = np.arange(N)
k = n.reshape((N, 1))
M = np.exp(-2j * np.pi * k * n / N)
return np.dot(M, x)
输出:
True
我们可以通过和numpy内置FFT函数比较再次检查这个计算结果
x = np.random.random(1024)
np.allclose(DFT_slow(x), np.fft.fft(x))
附:
def IDFT_slow(x):
"""Compute the inverse discrete Fourier Transform of the 1D array x"""
x = np.asarray(x, dtype=float)
N = x.shape[0]
n = np.arange(N)
k = n.reshape((N, 1))
M = np.exp(2j * np.pi * k * n / N)
return np.dot(x, M) / N
为了确认我们的算法有多么的慢,我们可以比较这两种方法的执行时间:
%timeit DFT_slow(x)
%timeit np.fft.fft(x)
输出:
10 loops, best of 3: 75.4 ms per loop
10000 loops, best of 3: 25.5 µs per loop
我们的算法慢将近1000倍,但这完全在我们的预料之中,因为这算法简单得过分。但这不是它最糟糕的表现,对于一个长度为N的输入向量,FFT算法的复杂度规模为
所以FFT是如何加速的呢?答案就藏在对称性的利用中。
利用一个问题的对称性是算法工程师腰带上最重要的工具之一。如果你能证明问题的一部分只和另一部分有关的话,你就可以只计算子结果一次从而减少计算开销。Cooley和Tukey在推导出FFT的过程中正使用了这种方法。
我们先开始计算以下
其中我们使用恒等式
最后一行展示出DFT一个很好的对称性:
Cooley和Tukey表明可以将DFT计算分成两个较小的部分。根据DFT的定义有:
我们将单个离散傅立叶变换分成两个项,它们本身看起来非常类似于较小的离散傅立叶变换,一个在奇数值上,一个在偶数值上。然而,到目前为止,我们还没有保存任何计算周期。每项包含
FFT的技巧就是利用每一项的对称性。因为k的范围是
但我们不能止步于此,只要我们更小的傅里叶变换有一个偶数值的M,我们就可以再次继续这个分治过程,每次减半计算量,直到我们的数组小到分治策略不再能带来性能提升。在渐进的极限中,这个递归方法复杂度规模为
这个递归的算法可以用python非常快的实现出来,当子问题足够小时就切回我们慢速版本的DFT代码:
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 <= 32: # this cutoff should be optimized
return DFT_slow(x)
else:
X_even = FFT(x[::2])
X_odd = FFT(x[1::2])
factor = np.exp(-2j * np.pi * np.arange(N) / N)
return np.concatenate([X_even + factor[:N / 2] * X_odd,
X_even + factor[N / 2:] * X_odd])
在这里,检查一下我们的算法的结果是否正确:
x = np.random.random(1024)
np.allclose(FFT(x), np.fft.fft(x))
输出:
True
同时我们对这个算法和慢速版本的算法进行运行速度上的比较:
%timeit DFT_slow(x)
%timeit FFT(x)
%timeit np.fft.fft(x)
输出:
10 loops, best of 3: 77.6 ms per loop
100 loops, best of 3: 4.07 ms per loop
10000 loops, best of 3: 24.7 µs per loop
我们的算法比naive版本的算法快出了一个数量级!更重要的是,我们递归算法复杂度渐进于
请注意,我们仍然没有Numpy内置FFT算法速度快,这是可想而知。Numpy使用的FFTPACK中的fft
算法是Fortran实现的,且经过了多年的调整和优化。此外,我们的NumPy解决方案涉及Python栈递归和许多临时数组的分配,这增加了大量的计算时间。
使用Python / NumPy时加速代码的一个好方法是尽可能对重复计算进行向量化。我们可以这样做,并在此过程中删除我们的递归函数调用,并使我们的Python FFT更高效。
请注意,在上面的递归的FFT实现中,在最深的递归调用中,都是相同的
def FFT_vectorized(x):
"""A vectorized, non-recursive version of the Cooley-Tukey FFT"""
x = np.asarray(x, dtype=float)
N = x.shape[0]
if np.log2(N) % 1 > 0:
raise ValueError("size of x must be a power of 2")
# N_min here is equivalent to the stopping condition above,
# and should be a power of 2
N_min = min(N, 32)
# Perform an O[N^2] DFT on all length-N_min sub-problems at once
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)))
# build-up each level of the recursive calculation all at once
while X.shape[0] < N:
X_even = X[:, :X.shape[1] / 2]
X_odd = X[:, X.shape[1] / 2:]
factor = np.exp(-1j * np.pi * np.arange(X.shape[0])
/ X.shape[0])[:, None]
X = np.vstack([X_even + factor * X_odd,
X_even - factor * X_odd])
return X.ravel()
虽然算法略微有些晦涩,但它只是重新排列递归版本中使用的操作,除了一个例外:我们在factor
的计算中利用了对称性,仅构造一半的数组。同样,我们将验证函数的正确性:
x = np.random.random(1024)
np.allclose(FFT_vectorized(x), np.fft.fft(x))
输出:
True
因为我们的算法变得更高效了,我们可以使用更大的数组来比较耗时,这里不列出DFT_slow
:
x = np.random.random(1024 * 16)
%timeit FFT(x)
%timeit FFT_vectorized(x)
%timeit np.fft.fft(x)
输出:
10 loops, best of 3: 72.8 ms per loop
100 loops, best of 3: 4.11 ms per loop
1000 loops, best of 3: 505 µs per loop
我们又将算法的速度提升了一个数量级!现在只比FFTPACK
慢10倍,而我们仅仅使用了几十行纯Python+Numpy。即使性能上弱于FFTPACK,但在可读性上远超FFTPACK
源码。
那么FFTPACK如何实现最后一点的加速呢?好吧,主要只是一个详细的簿记问题。FFTPACK花费大量时间确保重用任何可重用的子计算。我们的Numpy版本仍然涉及过多的内存分配和复制; 在像Fortran这样的语言中,它更容易控制并最大限度地减少内存使用。此外,Cooley-Tukey算法可推广为使用不是二分的切分策略(我们在这里实现的算法称为基-2 Cooley-Tukey FFT)。此外,可以使用其他更复杂的FFT算法,包括基于卷积的基本上不同的方法(参见例如Bluestein算法和Rader算法)。即使数组不是2的幂上,上述扩展和技术的组合也可以实现非常快的FFT。
虽然纯Python函数在实践中可能没用,但我希望它们能够让大家直观地了解基于FFT的数据分析背后发生的事情。作为数据科学家,我们可以使用更具算法头脑的同行构建的黑盒实现的基本工具,但我坚信对于我们进行数据处理的算法理解越深入,就越能发挥出它们的强大威力。