Solve Discrete Poisson Equation using Fast Fourier Transform (FFT)

1. Model

原文參考berkeley CS267

\(\operatorname{div}{G}\)為已知，目標是求\(I\)，滿足以下的二維Discrete Poisson Equation

\[
\begin{align}
\nabla^2{I}&=\operatorname{div}{G}\\
TI+IT&=\operatorname{div}{G} \label{eq:model1}
\end{align}
\]

其中\(T,I\)和\(\operatorname{div}{G}\)都是\(n\times{n}\)的矩陣，\(TI\)和\(IT\)分別代表對x和y方向做 laplacian，\(T\)是一個tridiagonal matrix，\(T=tridiag(1, -2, 1)\)

\[T=\begin{bmatrix}-2&1&0&0&\cdots&0\\1&-2&1&0&\cdots&0\\0&1&-2&1&\cdots&0\\\vdots&\vdots&\vdots&\vdots&\ddots&\vdots\\0&\cdots&1&-2&1&0\\0&\cdots&0&1&-2&1\\0&\cdots&0&0&1&-2\end{bmatrix}\]

2. Equations

並且\(T\)是一個non-singular matrix，做eigen decomposition，得到\(Ax=\lambda{x}\)的形式，其中\(Q\)的每一個column vector是為eigen vectors，而\(\lambda\)為一個對應的diagonal eigen value matrix

\[
\begin{align}
TQ&=\lambda{Q} \notag \\
T&=\lambda{Q}Q^{-1} \label{eq:T}
\end{align}
\]

將\(\eqref{eq:T}\)帶入\(\eqref{eq:model1}\)

\[
\begin{align}
(\lambda{Q}Q^{-1})I+I(\lambda{Q}Q^{-1})&=\operatorname{div}{G} \label{eq:e5}
\end{align}
\]

將\(\eqref{eq:e5}\)左乘\(Q^{-1}\)，右乘\(Q\)

\[
\begin{align}
\lambda{Q^{-1}}IQ+{Q^{-1}}IQ{\lambda}=Q^{-1}\operatorname{div}{G}{Q} \label{eq:e6}
\end{align}
\]

我們令\(\overline{x}=Q^{-1}xQ\)，將\(\eqref{eq:e6}\)改寫成易讀的形式

\[
\begin{align}
\lambda{\overline{I}}+{\overline{I}}\lambda&=\overline{\operatorname{div}{G}} \\
\lambda(j,j){\overline{I}(j,k)}+{\overline{I}(j,k)}\lambda{(k,k)}&=\overline{\operatorname{div}{G}}(j,k)
\end{align}
\]

因此經過簡單計算可得到每一項\(\overline{I}(j,k)\)

\[
\begin{align}
\overline{I}(j,k)=\frac{\overline{\operatorname{div}{G}}(j,k)}{\lambda{(j,j)}+\lambda{(k,k)}}
\end{align}
\]

最後反推求得\(I\)

\[
\begin{align}
I=Q\overline{I}Q^{-1}
\end{align}
\]

3. Algorithm

\[
\begin{align*}
&\overline{\operatorname{div}{G}}=Q^{-1}\operatorname{div}{G}Q \\
&\overline{I}(j,k)=\frac{\overline{\operatorname{div}{G}}(j,k)}{\lambda{(j,j)}+\lambda{(k,k)}} \\
&I=Q\overline{I}Q^{-1}
\end{align*}
\]

這裡有一個特性，因為\(T\)是一個symmetric matrix，所以其eigen vector \(Q^T=Q^{-1}=Q\)

簡化過的algorithm如下

\[
\begin{align*}
&\overline{\operatorname{div}{G}}=Q\operatorname{div}{G}Q \\
&\overline{I}(j,k)=\frac{\overline{\operatorname{div}{G}}(j,k)}{\lambda{(j,j)}+\lambda{(k,k)}} \\
&I=Q\overline{I}Q
\end{align*}
\]

其中最大的overhead就是和Q的矩陣乘法，由於\(Q\)是一個巨大的sparse matrix，僅在一行或一列中有三個非零項，直接儲存或是計算\(Q\)的矩陣乘法都是非常沒有效率且浪費資源的，我們勢必會找出\(Q\)的一些特性和規則，藉此加速該運算。

4. Multiplication with \(Q\)

\[TQ=\lambda{Q}\]

如果很用力的把他們每項展開得到下式

\[
\lambda(j,j)=2\left(1-\cos{\left(\frac{j\cdot\pi}{n+1}\right)}\right) \\
Q(j,k)=Q(k,j)=\sqrt{\frac{2}{n+1}}\sin\left(\frac{j\cdot{k}\cdot{\pi}}{n+1}\right)
\]

5. Accelerate using Fast Fourier Transform

因為\(Q\)是一個\(\sin\)的週期訊號，當然可以用FFT來加速計算，然後只取imaginary part coefficient。這邊注意\(Q\)的basis週期是\(\pi\)但ordinary FFT routine 中的定義是\(2\pi\)；若要用FFT加速需經過\(O(n)\)個zero padding的pre-processing來填補為兩倍週期，同理post-processing也應用在inverse FFT，最後再clipping原始週期片段。

左乘\(Q\)和右乘\(Q\)等價於\(x,y\)兩個方向的DST，也就是\(\overline{x}=QxQ=DST2D(x)\)，完整的加速演算法如下：

1. Zero padding \(\operatorname{div}{G}\) , double size in each direction x, y
2. Compute \(\overline{\operatorname{div}{G}}\) using 2D FFT, and get only the imaginary part coefficient.
3. Compute \(\overline{I}(j,k)=\frac{\overline{\operatorname{div}{G}}(j,k)}{\lambda{(j,j)}+\lambda{(k,k)}}\)
4. Compute \(I\) using 2D FFT, and get only the imaginary part coefficient.
5. Finally, crop the right sub-image of \(I\).

6. Discussion

在實作上，如果你的basis是1D做兩次的，建議一次僅填補一個dimension的zero padding，可以加速約\(\frac{1}{4}\)

有讀者一定會注意到，這邊FFT的cosine項都浪費掉了，如果將pre-processing和post-processing改成一個phase-shift的轉換，就可以僅用現有的DCT加速(比較多library提供DCT加速)。或是FFTW有提供僅做DST的選項。

但是筆者不是很喜歡這個推導方法，有點臭長，需要太多瑣碎的pre-knowledge(matrix, eigen-decomposition,frequency domain basis)。我喜歡從deconvolution的角度來看這個問題，下次篇將會帶領讀者們從Point Spread Function的基礎重新理解Poisson solver.