FFT 一瞥(1)

~11 min read
A journey from linear algebra to fast Fourier transform
tags: CodingMath

从线性变换到快速傅立叶变换(FFT) (1)

FFT(快速傅立叶变换) 是著名的算法, 将离散傅立叶变换的计算复杂度从 O(N2)\mathcal{O}(N^2) 变为了 O(NlogN)\mathcal{O}(N\log N) , 为了了解这一精妙算法, 以及其背后的对称性, 我们将从卷积出发, 从线性代数的视角看傅立叶变换是怎么自然出现, 以及其进一步利用群对称性的本质推广.

卷积是数学, 物理, 计算机等领域的常见计算需求, 在连续版本, 一维卷积的计算为

(fg)(x)=f(xy)g(y)dy\begin{equation} (f*g)(x)=\int_{-\infty}^{\infty}f(x-y)g(y)\mathrm{d}y \end{equation}

在算法上, 我们将函数变为一个向量, 实际上需要计算的是一个离散卷积

(fg)k=m=0n1fkmgm\begin{equation} (f*g)_k=\sum_{m=0}^{n-1} f_{k-m} g_{m} \end{equation}

我们将从线性代数的视角审视这一变换, 固定卷积函数 ff , 卷积实际上在做的是将向量 gg 变换为卷积向量 (fg)(f*g), 这满足线性变换的定义, 故这是一个在函数空间良好定义的线性变换, 我们可以使用研究线性空间的方法来研究这个变换

f,gf,g 作为序列都是在 modn\mathrm{mod} \,n 的意义上取

卷积作为线性变换

为了具体地研究卷积这个线性变换, 我们需要选择一组基, 不妨就直接取上面的基底(实际上就是标准正交基). 变换矩阵为

F=[f0fn1f2f1f1f0fn1f2f1f0fn2fn1fn1fn2f1f0]\begin{equation} F=\begin{bmatrix}f_{0}&f_{n-1}&\cdot&f_{2}&f_{1}\\f_{1}&f_{0}&f_{n-1}&\cdot&f_{2}\\\cdot&f_{1}&f_{0}&\cdot&\cdot\\f_{n-2}&\cdot&\cdot&\cdot&f_{n-1}\\f_{n-1}&f_{n-2}&\cdot&f_{1}&f_{0}\end{bmatrix} \end{equation}

我们发现非常周期的结构, 如果提取出 ff 的系数的话, 前三项分别是

f0×[e1e2en1en]f1×[e2e3ene1]f2×[e3e4e1e2]\begin{align} &f_0 \times [e_1\,e_2\,\cdots\,e_{n-1}\,e_{n}]\\ &f_1 \times [e_2\,e_3\,\cdots\, e_{n}\, e_{1}] \\ &f_2 \times [e_3\,e_4\,\cdots\, e_{1}\, e_{2}] \end{align}

这里 eie_iRn\mathbb{R}^n 上的正交基, 可以看到每部分实际上都是同一个矩阵(循环矩阵)的反复作用, 记为 PP ,那么实际上我们有 FF 的更简单结果

F=f0I+f1P+f2P2++fn1Pn1,\begin{equation} F=f_0I+f_{1}P+f_{2}P^2+\cdots+f_{n-1}P^{n-1}, \end{equation}

一般的卷积作用需要不断计算矩阵乘法, 如果我们能实现谱分解, 就可以让矩阵乘法变成数乘, 所以我们需要寻找 PP 的谱分解, 不难发现 PP 是一个 unitary matrix , 这有相当好的性质: 对于 unitary matrix , 我们知道它的特征值是单位根, 一共 n 个, 现在计算对应特征值的特征向量, 考虑单位根 ωk\omega^k , 我们有

Pχ=ωkχχj+1=ωkχj(modn)\begin{equation} P \chi = \omega^k \chi \quad\Longrightarrow\quad \chi_{j+1} = \omega^k \chi_j \quad (\mathrm{mod}\, n) \end{equation}

x0=1x_0 =1 , 我们发现

χk=(1,ωk,ω2k,,ω(n1)k)T,ω=e2πi/n\begin{equation} \chi_k = (1, \omega^k, \omega^{2k}, \ldots, \omega^{(n-1)k})^T, \quad \omega = e^{-2\pi i / n} \end{equation}

unitary matrix 的谱分解得到的本征向量自动是正交的, 我们取它们为一组基, 称为 Fourier basis, 把这些 basis 列在一起, 实际上就是标准基到 Fourier basis 的变换矩阵 WW:

Wkj=ωkj,ω=e2πi/n,k,j=0,1,,n1\begin{equation} W_{kj} = \omega^{kj}, \quad \omega = e^{-2\pi i / n}, \quad k,j = 0,1,\ldots,n-1 \end{equation}

W=[11111ωω2ωn11ω2ω4ω2(n1)1ωn1ω2(n1)ω(n1)2]\begin{equation} W = \begin{bmatrix} 1 & 1 & 1 & \cdots & 1 \\ 1 & \omega & \omega^2 & \cdots & \omega^{n-1} \\ 1 & \omega^2 & \omega^4 & \cdots & \omega^{2(n-1)} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \omega^{n-1} & \omega^{2(n-1)} & \cdots & \omega^{(n-1)^2} \end{bmatrix} \end{equation}

我们看 gg 在这一组基下的分量, 注意 WW 是 unitary 的 (W1=1nWW^{-1} = \frac{1}{n} W^*), 基变换为

g^=Wg,g^k=j=0n1gjωkj\begin{equation} \hat{g} = W g, \quad \hat{g}_k = \sum_{j=0}^{n-1} g_j\, \omega^{kj} \end{equation}

每个分量恰好对应傅立叶分量. 在这个基组下, PP 对每个基向量的作用是数乘:

Pχk=ωkχkP=W1diag(1,ω,ω2,,ωn1)W\begin{equation} P \chi_k = \omega^k \chi_k \quad\Longrightarrow\quad P = W^{-1} \,\mathrm{diag}(1, \omega, \omega^2, \ldots, \omega^{n-1})\, W \end{equation}

由于 F=m=0n1fmPmF = \sum_{m=0}^{n-1} f_m P^m, 代入谱分解即得 FF 的对角化:

F=W1diag(f^0,f^1,,f^n1)W,f^k=m=0n1fmωkm\begin{equation} F = W^{-1} \,\mathrm{diag}(\hat{f}_0, \hat{f}_1, \ldots, \hat{f}_{n-1})\, W, \quad \hat{f}_k = \sum_{m=0}^{n-1} f_m\, \omega^{km} \end{equation}

于是卷积在 Fourier 基下变为逐点乘积:

fg^=f^g^\begin{equation} \widehat{f * g} = \hat{f} \odot \hat{g} \end{equation}

这里 \odot 是 Hadamard 乘积 (逐元素乘). 这就是时域上的卷积定理——卷积在频域中化为数乘.

这就是离散版本的卷积定理, 可以发现, 为了计算卷积, 我们需要计算傅立叶变换, 而傅立叶变换需要计算变换矩阵, 而计算消耗是矩阵乘法 O(N2)\mathcal{O}(N^2) , 如何减小计算消耗就是我们下面需要讨论的问题.

分治: 幂次版本

DFT 的计算消耗是 N2N^2 量级, 为了减小计算消耗, 我们需要利用傅立叶核的对称性, 为了能看见这一点, 我们先介绍一个简单版本, 也是算法实现上更直观的版本 — Cooley-Tukey algorithm

如果 DFT 计算的阶数 NN 是偶数, 则 DFT 计算可以天然区分奇偶:

Xk=n=0N1xnei2πkn/N=m=0N/21x2mei2πk(2m)/N+m=0N/21x2m+1ei2πk(2m+1)/N=m=0N/21x2mei2πkm/(N/2)+ei2πk/Nm=0N/21x2m+1ei2πkm/(N/2)\begin{equation} \begin{aligned} X_{k} &= \sum_{n=0}^{N-1} x_n \cdot e^{-i 2\pi k n / N} \\ &= \sum_{m=0}^{N/2-1} x_{2m} \cdot e^{-i\left.2\pi\right.k\left(2m\right)/N} + \sum_{m=0}^{N/2-1} x_{2m+1} \cdot e^{-i\left.2\pi\right.k\left(2m+1\right)/N} \\ &= \sum_{m=0}^{N/2-1} x_{2m} \cdot e^{-i\left.2\pi\right.k\left.m\right/\left(N/2\right)} + e^{-i\left.2\pi\right.k/N} \sum_{m=0}^{N/2-1} x_{2m+1} \cdot e^{-i\left.2\pi\right.k\left.m\right/\left(N/2\right.)} \end{aligned} \end{equation}

注意到 NN 阶的 DFT 计算被我们变成了两个 N/2N/2 阶的 DFT 计算, 如果 DFT 计算还有 2 的因子, 则我们可以一直做下去, 直到做到一个较小的 nn 上我们能够直接计算 DFT 得到, 然后再递归得到整个 DFT 分量

代码实现

代码实现是简单的, 我们先给出简单的 DFT 计算程序

def DFT(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)

利用 numpy 的广播机制和向量化, 可以直接计算出变换矩阵, 然后递归构造 FFT

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(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])

使用 concatenate 拼接两个小的 FFT 成分: 注意保持频率大小次序: 对原先的 k 频率成分, 合成来自两个更小的 k 成分, 注意到奇偶频率在 modN/2\mathrm{mod} \,N/2 的意义上使用, 而外部的 factor 使用正常的 k

检查这确实符合频率成分的排序: x_even 和 x_odd 只储存了各自奇偶模式的成分, 组合为 k=0,…,n-1 种模式, 比如 1 和 N/2 +1 使用了相同的系数, 只不过是 factor 不同

向量化版本

实际上使用会直接采用向量化计算, 一次性算完需要进行的所有 O(N2)\mathcal{O}(N^2) DFT 计算

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()

逻辑大概为: 每一次递归产生 2 个新的 DFT 计算, 故最后递归到 NminN_{\min} 时产生了 N/NminN / N_{\min}NminN_{\min} 大小的 DFT 计算: 将需要进行 DFT 计算的向量排成列矩阵, 由变换矩阵一次性完成 DFT 计算, 然后不断拼接得到最终的频域分量即可

当然注意到这里有个结构: 在递归算法中, 我们需要的 DFT 序列不是连续的, 而 reshape 得到的每个片段是连续的向量. 实际发生的是 reshape 的矩阵, 每个列向量有 N_min 个坐标, 恰好按行数的列向量就是递归算法计算的 DFT 列向量, 变换矩阵左作用到矩阵上完成逐列向量的 DFT 变换. 故最后的每个列向量是一个 DFT 变换后的频域分量, 然后进行向上递推, 列向量是按照频率大小奇偶奇偶排列的, 然后接下来进行 vstack : 前半个频率排在上面, 后半个频率排在下面, 实现频率谱长度 x 2, 直到最后生成一个列向量, 这就是我们要的 DFT 频域向量

看起来现在我们的算法局限在处理 n=2kn=2^k 以及其附近的成分了, 如果 nn 是偶数, 则我们可以开心地进行递归, 但若 nn 不含 2 的因子, 则上述算法看似失效了? 这需要我们重新审视上面的步骤

循环群结构的分解: 混合基算法

分治的对称性

FFT 的算法只需要让我们从一个 DFT 计算中分拆出更小的 DFT 计算, 依靠递归来减小计算成本, 奇偶性只是一个我们容易拆分出更小 DFT 成分的情况, 实际发生的事情我们可以用这样一个单位根分布来表示, 图上是一个 16 阶 DFT 计算示意图, 我们每个傅立叶基表示为这样的 16 顶点

16 点 DFT 的基2 FFT 分解示意图

我们从这样的 16 边形中找到了两个更小周期的多边形: 8 顶点图, 正正好对应我们的奇偶性得到了两个更小规模的 DFT 计算, 所以我们发现 FFT 拆分的真正图像

能否从一个 n 顶点多边形图中, 找到更小的顶点多边形, 即找到一个更小的周期模式

我们从下面的 6 顶点图更能看出这一点

6 点 DFT 的混合基分解示意图

可以看见, 我们自然地从 6 顶点 DFT 中获得了两个 3 顶点 DFT , 这暗示着我们背后的群结构: 循环群的子群分解. 事实上, 循环群 ZN\mathbb{Z}_NNN 的每个因子 dd 恰好有一个 dd 阶子群, 而这个子群也是循环群. 于是若 N=N1N2N=N_1 N_2, 我们有子群 ZN2ZN\mathbb{Z}_{N_2} \subset \mathbb{Z}_N, 商群为 ZN/ZN2ZN1\mathbb{Z}_N / \mathbb{Z}_{N_2} \cong \mathbb{Z}_{N_1}, FFT 的拆分正是利用了这一结构: 通过子群将 DFT 分解为更小的 DFT 计算.

gcd(N1,N2)=1\gcd(N_1,N_2)=1 时, 中国剩余定理给出更强的结论:

ZNZN1×ZN2\mathbb{Z}_N \cong \mathbb{Z}_{N_1}\times\mathbb{Z}_{N_2}

此时 DFT 矩阵可以直接分解为两个较小 DFT 的张量积, 这是混合基算法的最佳情况. 但即使不互质, 子群—商群结构已足以驱动 Cooley-Tukey 分解.

有多种方式可以证明它, 总之, 如果这个成立, 我们立刻有: 对于一个 NN 顶点 DFT , 若 N=N1N2N=N_1N_2 我们可以将其递归到 N1N_1N2N_2 顶点 DFT 计算, 或者交换两个因子, 不难发现, 在前面的 6 顶点 DFT 中, 我们也能将其变为三个 2 顶点 DFT 计算

这就是 FFT 的混合基算法, 但是非 22 的素因子 DFT 向量化比较复杂, 我们这里就不再尝试代码实现

可以看到, 我们从一个计算优化问题, 先看到了奇偶自然产生的对称性, 并由此看到了背后更深的对称性, 从这个对称性出发, 我们也可以揭示 DFT 本身的对称性, 并将其推广.

不过混合基算法也是有局限的, 如果 nn 是一个素数, 则我们仍然没有有效的分治策略, 下次我们会正式从群论角度处理这个问题, 并由此引出基于循环群本身的广义傅立叶变换.

参考资料

← Blog