How to use matlab command 'fft' to solve Ax=b arising from Poisson equation?

Question

I want to ask a question about fast solver to the Poisson equation with Homogenous boundary conditions as follows: $$-\Delta u = f.$$ After centered difference using $n+2$ equidistance points in all directions, we get a SPD system: $$ Ax=b. $$ Let $T = \begin{bmatrix} 2 & -1 & & \\ -1 & 2 & {\ddots} & \\ & {\ddots} & {\ddots} & -1 \\ & & -1 & 2 \\ \end{bmatrix}\in \mathbb{R}^{n\times n}$, then we have

1. in 1D case, $A=T\in \mathbb{R}^{n\times n}$;
2. in 2D case, $A = T\otimes I +I\otimes T\in \mathbb{R}^{n^2\times n^2}$;
3. in 3D case, $A = T\otimes I\otimes I+I\otimes T\otimes I+I\otimes I\otimes T\in \mathbb{R}^{n^3\times n^3},$ where $\otimes$ denotes the kroneck product.

I have known that in 1D and 2D, matlab command A\b is very fast and I have tested in matlab. But in 3D 'A\b' tends to be slow for large sparse matrix $A$. I have heard that in matlab, there is a function 'fft' (fast fourier transformation) can solve Ax=b quickly, but I can not understand that. Can someone give me the matlab impelementation codes that I can easily understand how matlab use 'fft' to solve this Ax=b quickly? Thanks very much.

Below is my matlab code for n=50 such that A\b fails in my matlab2018b (CPU 8G memory):

clc;clear;rng(0);
n=50;
T = spdiags(ones(n,1).*[-1 2 -1], [-1 0 1],n,n);
I = speye(n);
A = kron(kron(T,I),I)+kron(I,kron(T,I))+kron(I,kron(I,T));
N = length(A);
b = rand(N,1);
tic
A\b;
toc

Answer 1

Let $T \in \mathbb{R}^{n \times n}$ be defined as in your question, $$ T = \begin{pmatrix} 2 & -1 & & \\ -1 & 2 & {\ddots} & \\ & {\ddots} & {\ddots} & -1 \\ & & -1 & 2 \\ \end{pmatrix}. $$ Then, it is well known that the eigenvalue decomposition of $T$ is given by $$ T = Z \Lambda Z, $$ where $\Lambda$ is diagonal with $$ \Lambda_{jj} = 2\left(1-\cos\left( \frac{\pi j}{n+1} \right)\right) $$ and $Z$ is given by $$ Z_{jk} = \sqrt{\frac{2}{n+1}}\sin\left(\frac{\pi j k}{n+1}\right). $$ Notice that $Z = Z^T = Z^{-1}$.

Then, e.g. for the two-dimensional Laplacian, we have \begin{align*} A &= I \otimes T + T \otimes I \\ & = I \otimes (Z\Lambda Z) + (Z\Lambda Z) \otimes I \\ &= (Z \otimes Z) (I \otimes \Lambda + \Lambda \otimes I) (Z \otimes Z). \end{align*}

So, we can solve $A u = b$ according to the following algorithm:

1. Let $b' = (Z \otimes Z) b$.
2. Solve $u' = (I \otimes \Lambda + \Lambda \otimes I)^{-1} b'$ (notice that this matrix is diagonal, and the values of $\Lambda$ are known using the above formula).
3. Compute $u = (Z \otimes Z) u'$.

The algorithm for the $d$-dimensional Laplacian is obtained analogously.

FFT Poisson solvers are based on using the FFT to quickly compute products of the form $(Z \otimes \cdots \otimes Z) b$. The idea is as follows.

The discrete Fourier transform (DFT) of a vector $v \in \mathbb{R}^m$ is given by the matrix-vector product $\Phi v$ where the matrix $\Phi \in \mathbb{R}^{m \times m}$ is defined by $$ \Phi_{jk} = \exp\left( \frac{- 2 \pi i j k }{m} \right). $$

Then, notice that multiplying a vector $b \in \mathbb{R}^n$ by $Z$ to compute the product $Z b$ corresponds to taking a subvector of the imaginary part of the DFT of a padded vector. In particular, define $\tilde{b} = (0, b, 0, \ldots, 0) \in \mathbb{R}^{2n+2}$ (i.e. padding $b$ on the left with one zero, and on the right with $n+1$ zeros). Then, $Zb$ is equal to $\sqrt{2/(n+1)}$ times the second through $(n+1)$st entries of the imaginary part of the size-$(2n+2)$ DFT of $\tilde{b}$. We use the FFT to compute the DFT of $\tilde{b}$ efficiently.

This idea generalizes to arbitrary dimensions, by considering the vector $b$ to be a $d$-tensor, and then constructing $\tilde{b}$ by padding with zeros along each dimension, and then applying the FFT in succession along each dimension.

See below for a proof-of-concept for the two-dimensional Laplacian. It should be possible to extend this to 3D with a bit of book-keeping.

n = 50;
T = spdiags(ones(n,1).*[-1 2 -1],[-1 0 1],n,n);
I = speye(n);
A = kron(T,I) + kron(I,T);

lambda = zeros(n,1);
for i=1:n
   lambda(i) = 2*(1-cos(pi*i/(n+1)));
end
L = lambda + lambda';

% random right-hand side
b = rand(n*n,1);
B = reshape(b,n,n);

% pad with zeros
B_ext = [zeros(1,n) ; B ; zeros(n+1,n)];
B_ext = [zeros(2*n+2,1) B_ext zeros(2*n+2, n+1)];

% compute (Z \otimes Z) B
c = -sqrt(2/(n+1));
B_prime = imag(fft(B_ext, [], 2));
B_prime = imag(fft(B_prime(:,2:n+1), [], 1));
B_prime = c^2*B_prime(2:n+1,:);

% solve diagonal system
U_prime = B_prime ./ L;

% pad with zeros
U_prime_ext = [zeros(1,n) ; U_prime ; zeros(n+1,n)];
U_prime_ext = [zeros(2*n+2,1) U_prime_ext zeros(2*n+2, n+1)];

% compute (Z \otimes Z) U_prime
U = imag(fft(U_prime_ext, [], 2));
U = imag(fft(U(:,2:n+1), [], 1));
U = c^2*U(2:n+1,:);