Python计算傅里叶变换

技术背景

傅里叶变换在几乎所有计算相关领域都有可能被使用到，例如通信领域的滤波、材料领域的晶格倒易空间计算还有分子动力学中的倒易力场能量项等等。最简单的例子来说，计算周期性盒子的电势能$k\sum_i\frac{q_i}{r_i}$本身就是一个类似于调和级数的形式，很难求得精确解。但是在Edward求和方法中使用傅里叶变换，可以做到在倒易空间进行能量计算，可以逼近精确解。本文主要介绍傅里叶变换的原理及相应的Python代码实现。

DFT原理

DFT计算的本质是一个矩阵运算：

\[y_k=\sum_{n=0}^{N-1}x_ne^{-j\frac{2\pi nk}{N}},0\leq k\leq N-1 \]

\[x_n=\frac{1}{N}\sum_{k=0}^{N-1}y_ke^{j\frac{2\pi nk}{N}},0\leq n\leq N-1 \]

如果写成一个矩阵的形式，那就是：

\[\left[\begin{matrix} y_1\\ y_2\\ y_3\\ ...\\ y_{N-1} \end{matrix}\right]=\left[\begin{matrix} 1&&1&&1&&...&&1\\ 1&&e^{-j\frac{2\pi}{N}\cdot1}&&e^{-j\frac{2\pi}{N}\cdot2}&&...&&e^{-j\frac{2\pi}{N}\cdot(N-1)}\\ 1&&e^{-j\frac{2\pi}{N}\cdot2}&&e^{-j\frac{2\pi}{N}\cdot4}&&...&&e^{-j\frac{2\pi}{N}\cdot2(N-1)}\\ ...&&...&&...&&...&&...\\ 1&&e^{-j\frac{2\pi}{N}\cdot(N-1)}&&e^{-j\frac{2\pi}{N}\cdot2(N-1)}&&...&&e^{-j\frac{2\pi}{N}\cdot(N-1)(N-1)} \end{matrix}\right]\left[ \begin{matrix} x_1\\ x_2\\ x_3\\ ...\\ x_{N-1} \end{matrix}\right] \]

类似的，逆傅里叶变换的矩阵形式为：

\[\left[ \begin{matrix} x_1\\ x_2\\ x_3\\ ...\\ x_{N-1} \end{matrix}\right] =\left[\begin{matrix} 1&&1&&1&&...&&1\\ 1&&e^{j\frac{2\pi}{N}\cdot1}&&e^{j\frac{2\pi}{N}\cdot2}&&...&&e^{j\frac{2\pi}{N}\cdot(N-1)}\\ 1&&e^{j\frac{2\pi}{N}\cdot2}&&e^{j\frac{2\pi}{N}\cdot4}&&...&&e^{j\frac{2\pi}{N}\cdot2(N-1)}\\ ...&&...&&...&&...&&...\\ 1&&e^{j\frac{2\pi}{N}\cdot(N-1)}&&e^{j\frac{2\pi}{N}\cdot2(N-1)}&&...&&e^{j\frac{2\pi}{N}\cdot(N-1)(N-1)} \end{matrix}\right] \left[\begin{matrix} y_1\\ y_2\\ y_3\\ ...\\ y_{N-1} \end{matrix}\right] \]

如果记参数$W_{N,n,k}=e^{-j\frac{2\pi}{N}nk}$，则其共轭$W_{N,n,k}^*=e^{j\frac{2\pi}{N}nk}$是逆傅里叶变换的参数。而且根据复变函数的性质，该参数具有周期性：$W_{N,n+N,k}=W_{N,n,k+N}=W_{N,n,k}$，共轭参数同理。最后还有一个非常重要的性质：$W_{N/m,n/m,k}=W_{N/m,n,k/m}=W_{N,n,k}$，根据这个特性，可以将大规模的运算变成小范围的计算。在不考虑这些参数特性的情况下，我们可以使用Python做一个初步的DFT简单实现。

初步Python实现

这里没有做任何的优化，仅仅是一个示例：

import numpy as np

def dft(x):
    y = np.zeros_like(x, dtype=np.complex64)
    N = x.shape[0]
    for k in range(N):
        y[k] = np.sum(x * np.exp(-1j*2*np.pi*k*np.arange(N)/N))
    return y

def idft(y):
    x = np.zeros_like(y, dtype=np.float32)
    N = y.shape[0]
    for n in range(N):
        x[n] = np.real(np.sum(y * np.exp(1j*2*np.pi*n*np.arange(N)/N)) / N)
    return x

N = 128
x = np.random.random(N).astype(np.float32)
y0 = dft(x)
y1 = np.fft.fft(x)
print (np.allclose(y0, y1))

yr = np.random.random(N).astype(np.float32)
yi = np.random.random(N).astype(np.float32)
y = yr + 1j*yi
x0 = idft(y)
x1 = np.fft.ifft(y).real
print (np.allclose(x0, x1))
# True
# True

输出的两个结果都是True，也就说明这个计算结果是没问题的。

FFT快速傅里叶变换

首先我们整理一下所有参数相关的优化点：

\[\left\{ \begin{matrix} W_{N,n+N,k}=W_{N,n,k+N}=W_{N,n,k}\\ W_{N/m,n/m,k}=W_{N/m,n,k/m}=W_{N,n,k}\\ W_{N,\beta,\frac{N}{2\beta}}=W^*_{N,\beta,\frac{N}{2\beta}}=-1\Rightarrow W_{N,n,k}\cdot W_{N,\beta,\frac{N}{2\beta}}=-W_{N,n,k}\\ W_{N,N-n,k}=W_{N,N,k}\cdot W_{N,-n,k}=W_{N,-n,k}=W_{N,n,N-k} \end{matrix} \right. \]

此时如果我们把原始的输入$x_n$拆分为奇偶两组（如果总数N不是偶数，一般可以对输入数组做padding）：

\[\left\{ \begin{matrix} x_{2r}\\ x_{2r+1} \end{matrix},0\leq r\leq \frac{N}{2}-1 \right. \]

则有：

\[y_k=\sum_{n=0}^{N-1}x_ne^{-j\frac{2\pi nk}{N}}=\sum_{r=0}^{\frac{N}{2}-1}x_{2r}W_{N,2r,k}+\sum_{r=0}^{\frac{N}{2}-1}x_{2r+1}W_{N,2r+1,k} \]

如果我们把$x_{2r}$和$x_{2r+1}$看作是两个独立的输入数据，那么上述分解可以进一步优化：

\[\begin{align*} y_k&=\sum_{r=0}^{\frac{N}{2}-1}x_{2r}W_{N,2r,k}+\sum_{r=0}^{\frac{N}{2}-1}x_{2r+1}W_{N,2r+1,k}\\ &=\sum_{r=0}^{\frac{N}{2}-1}x^{(odd)}_{r'}W_{\frac{N}{2},r',k}+W_{N,1,k}\sum_{r=0}^{\frac{N}{2}-1}x^{(even)}_{r'}W_{\frac{N}{2},r',k}\\ &=y^{(odd)}_k+W_{N,1,k}y^{(even)}_k \end{align*} \]

同理可以得到：

\[\begin{align*} y_{k+\frac{N}{2}}&=y^{(odd)}_{k+\frac{N}{2}}+W_{N,1,k+\frac{N}{2}}y^{(even)}_{k+\frac{N}{2}}\\ &=y^{(odd)}_{k+\frac{N}{2}}-W_{N,1,k}y^{(even)}_{k+\frac{N}{2}} \end{align*} \]

这就是所谓的蝶形运算（图像来自于参考链接）：

这个运算式的意义在于，假如我们原本做一个$2^N$点数据的傅里叶变换，使用原始的DFT运算我们需要做$2^{2N}$次乘法和$2^N(2^N-1)$次加法，但是这种方法可以把计算量缩减到$2\cdot2^N+2^{\frac{N}{2}}$次乘法和$2^{\frac{N}{2}}(2^\frac{N}{2}-1)$次加法。做一次分解，就把复杂度从$O(2^{2N})$降到了$O(2^N)$（注意：这里的$N$跟前面用到的数据点总数不是一个含义，这里的$N$指代数据点总数是2的整数次方，只是两者的表述习惯都常用$N$）。相关代码实现如下：

import numpy as np

def dft(x):
    y = np.zeros_like(x, dtype=np.complex64)
    N = x.shape[0]
    for k in range(N):
        y[k] = np.sum(x * np.exp(-1j*2*np.pi*k*np.arange(N)/N))
    return y

def dft2(x):
    y = np.zeros_like(x, dtype=np.complex64)
    N = x.shape[0]
    for k in range(N//2):
        c1 = np.exp(-1j*2*np.pi*k*np.arange(N//2)/(N//2))
        c2 = np.exp(-1j*2*np.pi*k/N)
        y1 = np.sum(x[::2] * c1)
        y2 = np.sum(x[1::2] * c1)
        y[k] = y1 + c2 * y2
        y[k+N//2] = y1 - c2 * y2
    return y

N = 128
x = np.random.random(N).astype(np.float32)
y0 = dft2(x)
y1 = np.fft.fft(x)
print (np.allclose(y0, y1))
# True

运行输出为True，表示计算结果一致。需要注意的是，这里的代码未考虑padding问题，不能作为正式的代码实现，仅仅是一个算法演示。既然能够分割一次，那么就可以分割多次，直到无法分割为止，或者分割到一个指定的参数为止。这也就是多重蝶形运算的原理：

简单一点可以使用递归的方式进行计算：

import numpy as np

def dft(x):
    y = np.zeros_like(x, dtype=np.complex64)
    N = x.shape[0]
    for k in range(N):
        y[k] = np.sum(x * np.exp(-1j*2*np.pi*k*np.arange(N)/N))
    return y

def dftn(x, N_cut=2):
    y = np.zeros_like(x, dtype=np.complex64)
    N = x.shape[0]
    if N > N_cut:
        y1 = dftn(x[::2])
        y2 = dftn(x[1::2])
    else:
        return dft(x)
    for k in range(N//2):
        c2 = np.exp(-1j*2*np.pi*k/N)
        y[k] = y1[k] + c2 * y2[k]
        y[k+N//2] = y1[k] - c2 * y2[k]
    return y

N = 1024
x = np.random.random(N).astype(np.float32)
y0 = dftn(x)
y1 = np.fft.fft(x)
print (np.allclose(y0, y1))
# True

这里的实现使用递归的方法，结合了前面实现的DFT算法和蝶形运算方法，得到的结果也是正确的。这里使用的蝶形运算优化方法，就是FFT快速傅里叶变换的基本思路。

N点快速傅里叶变换

所谓的N点FFT，其实就是每次只取N个数据点执行傅里叶变换。那么取数据点的方式就有很多种了，例如只取前N个数据点，或者降采样之后再取前N个数据点，再就是加窗，在经过窗函数的运算后，对每个窗体内的数据点做傅里叶变换。最简单的方式就是矩形窗，常见的还有汉宁窗和汉明窗，这里不做详细分析。值得注意的是，如果使用降采样的方法，采样率需要遵循奈奎斯特采样定理，要大于两倍的target frequency。尤其对于周期性边界条件和远程相互作用的场景，高频区域的贡献是不可忽视的。

至于为什么不使用全域数据点的傅里叶变换，即使我们可以用快速傅里叶变换把计算复杂度缩减到$O(N\log N)$（这里的$N$是数据点数）的级别，对于那些大规模数据传输和计算的场景，也是不适用的，因此使用降低傅里叶变换点数的思路对于大多数的场景来说可以兼顾到性能与精确度。而窗体函数的出现，进一步优化了截断处数据泄露的问题。

总结概要

本文介绍了离散傅里叶变换和快速傅里叶变换的基本原理及其对应的Python代码实现，并将计算结果与numpy所集成的fft函数进行对比。其实现在FFT计算的成熟工具已经有很多了，不论是CPU上scipy的fft模块还是GPU上的cufft动态链接库，都有非常好的性能。但还是得真正去了解计算背后的原理，和相关的物理图像，才能更恰当的使用这个强大的工具。