Small box fast Fourier transform（SBFFT）算法是一种将空间分为小块并行进行傅里叶变换的算法，其优点在于相比传统的基于一维数据划分的FFT算法可以使用更多的核。本系列文章是学习这一算法在求解泊松方程相关代码的笔记，本文是其中的第一篇，主要参考了文献 [1]。
在DFT计算中，从电荷密度 \rho(\boldsymbol{r}) 计算库仑势 \phi(\boldsymbol{r}) 的过程，是通过在给定的边界条件下求解泊松方程
\nabla^2\phi(\boldsymbol{r})=-4\pi\rho(\boldsymbol{r})使用格林函数方法，可以将解写成
\phi(\boldsymbol{r})=\int V(\boldsymbol{r-r}_1)\rho(\boldsymbol{r}_1)\,\mathrm{d}^3\boldsymbol{r}_1在常用的周期性边界条件，这一积分是在有限的格子（通常是单胞）中进行的，而其中的核函数V(\boldsymbol{r-r}_1)=1/|\boldsymbol{r-r}_1|。注意到这一积分是一个卷积的形式，那么在对相应的函数做傅里叶变换f(\boldsymbol{G})=\mathscr{F}\{f(\boldsymbol{r})\}，得到在倒空间中的形式
\phi(\boldsymbol{G})=V(\boldsymbol{G})\rho(\boldsymbol{G})其中V(\boldsymbol{G})=4\pi/|\boldsymbol{G}|^2。因此库伦势\phi(\boldsymbol{r})可以通过一个傅里叶变换（\rho(\boldsymbol{G})= \mathscr{F}\{\rho(\boldsymbol{G})\}）和一个傅里叶逆变换 （ \phi(\boldsymbol{r})= \mathscr{F}^{-1}\{\phi(\boldsymbol{G})\} ）得到。
在目前常用的（并行版本）FFTW中，对于一个N_1\times N_2\times N_3的三维数组的傅里叶变换，能使用的核数最多为N_3（Fortran）或 N_1（C） 。对于一个大数组，这样的并行度显然还不够高。因此最好是能够实现一种分块的FFT算法，不仅能够使用更多的核，而且能够尽量降低全局通信。然而这里的核函数V(\boldsymbol{r-r}')是长程的，如果直接使用分块算法需要大量的全局通信，对此的解决思路是将这一核函数分解为短程项和长程项两部分
V(\boldsymbol{r-r}_1) = V_\text{S}(\boldsymbol{r-r}_1)+ V_\text{L}(\boldsymbol{r-r}_1)这里的关键在于选取合适的分解使得短程项V_\text{S}（在实空间）是短程的而长程项V_\text{L}是平滑的（因而在倒空间中是短程的），从而可以利用各自在不同空间的短程特性分别使用分块FFT算法进行处理
\begin{aligned} \phi_\text{S}(\boldsymbol{r})&=\int V_\text{S}(\boldsymbol{r-r}_1)\rho(\boldsymbol{r}_1)\,\mathrm{d}^3\boldsymbol{r}_1 \\ \phi_\text{L}(\boldsymbol{r})&=\int V_\text{L}(\boldsymbol{r-r}_1)\rho(\boldsymbol{r}_1)\,\mathrm{d}^3\boldsymbol{r}_1 \end{aligned}最后将得到的短程库仑势\phi_\text{S}(\boldsymbol{r})和长程库仑势\phi_\text{L}(\boldsymbol{r})相加即可得到库仑势\phi(\boldsymbol{r})。
1. 短程项的选取及其处理
由于核函数V无论在实空间还是倒空间中都是球对称的，因此可以把对其的分解简化为一个一维问题，其一维傅里叶变换的形式为
V(r)=\frac{1}{2\pi^2r}\int\sin(Gr)GV(G)\,\mathrm{d}G数值计算中使用离散的长度为N数组r_i=iN/L和G_j=j\pi/L，那么这一积分的离散形式为
V(r_i)=\sum_{j=1}^{N}\frac{jN}{2iL^3} \sin(ij\pi/N)V(G_j)= \sum_{j=1}^{N}C_{ij}V(G_j)对于给定的实空间和倒空间中的截断长度r_\text{c}和G_\text{c}，可以将V_\text{S}在实空间和V_\text{L}在倒空间中的短程特性写成
\begin{aligned} &V_\text{S}(r_i)=0, \forall i>r_\text{c}N/L\\ &V_\text{L}(G_j)=0, \forall j>G_\text{c}L/\pi \end{aligned}由于V_\text{S}+V_\text{L}=V，可以通过使得V_\text{S}(r)在 r_\text{c}以外的部分为0来进行求解，即最小化如下的量
\begin{aligned} O=&\sum_{i=r_\text{c}N/L}^N|V_\text{S}(r_i)|^2r_i^2\\=&\left(\frac{L}{N}\right)^2\sum_{i=r_\text{c}N/L}^Ni^2\left[\sum_{j=1}^{N}C_{ij}\frac{4\pi}{G_j^2}- \sum_{j=1}^{G_\text{c}L/\pi}C_{ij}V_\text{L}(G_j)\right]^2 \end{aligned}求解对应的线性方程组即可求得V_\text{L}(G_j)，从而可以将核函数分解为短程项V_\text{S}和长程项V_\text{L}，示例结果如下图所示。注意到这里得到的V_\text{S}(r)在r_\text{c}以外很小（<10^{-5}）并不严格为0，而V_\text{L}(G)在G_\text{c}以外严格为0，这是由求解方式决定的。
对于短程项部分的处理可以使用SBFFT算法，将体系在实空间中划分成一系列三维的小块（称为核心区域），对于每一个核心区域\Omega_\text{I}在各个方向向外扩展长度L_\text{c}（其中L_\text{c}>r_\text{c}）作为缓冲区，从而成为实际计算的小块\Omega_\text{II}。由于V_\text{S}在缓存区外为0，因此各个小块\Omega_\text{II}的FFT可以并行处理，仅需要与附近26个小块进行通信（交换缓冲区数据），避免了全局通信。这样的计算简记为DG-SBFFT，这里的DG表示原始的稠密密网格，与下文中的稀疏网格SG对应。
2. 长程项的处理
对于长程库仑势的计算是比较复杂的，这里处理的关键的在于利用长程项V_\text{L}(G)在G>G_\text{c}时为0的特性。由于计算中使用的最大的倒格矢G_\text{max}=N\pi/L大于G_\text{c}（可以通过选择合适的G_\text{c}实现），这表明可以用一个较稀疏的实空间网格（对应较小G'_\text{max}）来进行计算而并不会影响精度。为了方便起见，下面的讨论中分别用不加上标和加上标'区分原始稠密网格和新引入的稀疏网格上的变量，用不加上标和加上标^\text{s}区分全局和小块中网格的变量。在倒空间中对于长程项有如下关系
\begin{aligned} \phi_\text{L}(\boldsymbol{G})&=\rho(\boldsymbol{G})V_\text{L}(\boldsymbol{G})\\&=[\rho(\boldsymbol{G})m(\boldsymbol{G})]\cdot[V_\text{L}(\boldsymbol{G})/m(\boldsymbol{G})]\\&=\rho_\text{m}(\boldsymbol{G})\cdot[V_\text{L}(\boldsymbol{G})/m(\boldsymbol{G})] \end{aligned}这里的m(\boldsymbol{G})是一个（倒空间中）短程的mask函数，当G>G_\text{mc}时为0（这里G_\text{mc}>G_\text{c}）保证了V_\text{L}(\boldsymbol{G})/m(\boldsymbol{G})是定义良好的。将 \rho_\text{m}(\boldsymbol{G})=\rho(\boldsymbol{G})m(\boldsymbol{G})变换到实空间中有
\rho_\text{m}(\boldsymbol{G})=\int m(\boldsymbol{r-r}_1)\rho(\boldsymbol{r}_1)\,\mathrm{d}^3\boldsymbol{r}_1如果m(\boldsymbol{r})在实空间中也是是短程的，即r>r_\text{c}时其值为0，那么对\rho_\text{m}(\boldsymbol{r})可以采用上一小节中对于短程项相同的方法计算得到。这里引入\rho_\text{m}的目的是利用其在倒空间中的短程特点，对其及\phi_\text{L}的FFT可以在一个较为稀疏的全局网格上进行，从而大大降低了计算量。注意到这里对函数m的要求是在实空间和倒空间中都是短程的，与上一小节中分解核函数的要求相似，可以类似的方法得到。
因此，对于（稀疏网格上的）长程库仑势\phi'_\text{L}的计算可以分为三个步骤。第一步，用计算短程库仑势中得到各小块中的\rho^\text{s}(\boldsymbol{G}^\text{s})乘以mask函数m(\boldsymbol{G}^\text{s})得到\rho^\text{s}_\text{m}(\boldsymbol{G}^\text{s}) ，由于其在G^\text{s}>G_\text{mc}时为0，因此可以简化为一个G_\text{max}较小的\rho'^\text{s}_\text{m}(\boldsymbol{G}^\text{s})，对应实空间中一个较为稀疏的网格。对\rho'^\text{s}_\text{m}(\boldsymbol{G}^\text{s})做逆SBFFT计算，得到\rho'^\text{s}_\text{m}(\boldsymbol{r}'^\text{s})。由于这一变换是在稀疏网格上进行的，简记为SG-SBFFT。第二步，将各个小块中的\rho'^\text{s}_\text{m}(\boldsymbol{r}'^\text{s})拼接得到全局的\rho'_\text{m}(\boldsymbol{r}')，对其进行FFT计算即可得到\rho'_\text{m}(\boldsymbol{G})，这一变换简记为SG-FFT。第三步，用上一步得到的\rho'_\text{m}(\boldsymbol{G})乘以V_\text{L}(\boldsymbol{G})/m(\boldsymbol{G})得到\Phi'_\text{L}(\boldsymbol{G})，进行一次逆SG-FFT计算即可得到\Phi'_\text{L}(\boldsymbol{r}')。
最后需要将稀疏网格上的\Phi'_\text{L}(\boldsymbol{r}')变换为原始上网格\Phi_\text{L}(\boldsymbol{r})。这一问题与前面的稠密网格变换为稀疏网格是逆问题，因此同样可以通过SBFFT来解决。这一步同样需要引入一个在实空间和倒空间中均为短程的mask函数M，值得注意的是这里需要使用三个一维mask函数相乘的形式并将截断长度L_\text{s}/2（L_\text{c}/2为小块核心区域的长度，这是为了保证在边界上不存在跳变）。
3. 总结
根据上面的讨论，可以将整个计算过程的流程总结如下：
- 短程库仑势的计算
- \rho(\boldsymbol{r})\rightarrow\rho^\text{s}(\boldsymbol{r}^\text{s})（赋值）
- \rho^\text{s}(\boldsymbol{r}^\text{s})\rightarrow\rho^\text{s}(\boldsymbol{G}^\text{s})（DG-SBFFT）
- \phi^\text{s}_\text{S}(\boldsymbol{r}^\text{s})=V_\text{s}(\boldsymbol{G}^\text{s})\rho^\text{s}(\boldsymbol{G}^\text{s})
- \phi_\text{S}^\text{s}(\boldsymbol{G}^\text{s})\rightarrow\phi_\text{S}^\text{s}(\boldsymbol{r}^\text{s})（DG-SBFFT）
- \phi^\text{s}_\text{S}(\boldsymbol{r}^\text{s})\rightarrow\phi_\text{S}(\boldsymbol{r}) （拼合）
- 长程库仑势的计算
- \rho^\text{s}_\text{m}(\boldsymbol{G}^\text{s})=m(\boldsymbol{G}^\text{s})\rho^\text{s}(\boldsymbol{G}^\text{s})
- \rho^\text{s}_\text{m}(\boldsymbol{G}^\text{s})\rightarrow\rho'^\text{s}_\text{m}(\boldsymbol{G}^\text{s})（去掉高频项）
- \rho'^\text{s}_\text{m}(\boldsymbol{G}^\text{s})\rightarrow\rho'^\text{s}_\text{m}(\boldsymbol{r}'^\text{s})（SG-SBFFT）
- \rho'^\text{s}_\text{m}(\boldsymbol{r}'^\text{s})\rightarrow\rho'_\text{m}(\boldsymbol{r}')（拼合）
- \rho'_\text{m}(\boldsymbol{r}')\rightarrow\rho'_\text{m}(\boldsymbol{G})（SG-FFT）
- \phi'_\text{L}(\boldsymbol{G})=\rho'_\text{m}(\boldsymbol{G})\cdot[V_\text{L}(\boldsymbol{G})/m(\boldsymbol{G})]
- \phi'_\text{L}(\boldsymbol{G})\rightarrow\phi'_\text{L}(\boldsymbol{r}')（SG-FFT）
- \phi'_\text{L}(\boldsymbol{r}')\rightarrow\phi'^\text{s}_\text{L}(\boldsymbol{r}'^\text{s})（赋值）
- W'^\text{s}(\boldsymbol{r}'^\text{s})=M'(\boldsymbol{r}'^\text{s}) \phi'^\text{s}_\text{L}(\boldsymbol{r}'^\text{s})
- W'^\text{s}(\boldsymbol{r}'^\text{s})\rightarrow W'^\text{s}(\boldsymbol{G}^\text{s}) （SG-SBFFT）
- W'^\text{s}(\boldsymbol{G}^\text{s})\rightarrow W^\text{s}(\boldsymbol{G}^\text{s}) （高频部分补零）
- W^\text{s}(\boldsymbol{G}^\text{s})\rightarrow W^\text{s}(\boldsymbol{r}^\text{s}) （DG-SBFFT）
- \phi^\text{s}_\text{L}(\boldsymbol{r}^\text{s} )=W^\text{s} (\boldsymbol{r}^\text{s})/M(\boldsymbol{r}^\text{s} )
- \phi^\text{s}_\text{L}(\boldsymbol{r}^\text{s})\rightarrow\phi_\text{L}(\boldsymbol{r})（拼合）
- 短程项与长程项相加
- \phi(\boldsymbol{r})=\phi_\text{S}(\boldsymbol{r})+\phi_\text{L}(\boldsymbol{r})
可以看到这里将原来算法中需要的两个全局DG-FFT替换成了两个全局SG-FFT、三个DG-SBFFT以及两个SG-SBFFT。
参考文献
[1] “A small box Fast Fourier Transformation method for fast Poisson solutions in large system“.