Q3 Fast PDE solver: Poisson's equation¶
In HW2 we derived a numerical scheme to solve Poisson's equation in 1D:
$$
- \frac{d^2}{dx^2} u(x) = f(x) \qquad u(0) = u(1) = 0
$$
The discretized system was of the form:
$$
T u = h^2 f
$$
where:
$$
T = \begin{bmatrix}
2 & -1 & & & & \\
-1 & 2 & -1 & & & \\
& -1 & 2 & -1 & & \\
& & \ddots & \ddots & \ddots & \\
& & & -1 & 2 & -1 \\
& & & & -1 & 2 \\
\end{bmatrix},
f = \begin{bmatrix}
f(h) \\
f(2h) \\
\vdots \\
f(n h)
\end{bmatrix}
$$
T is tridiagonal, so we now know how to solve this system in $O(n)$ operations. All variations of the 1D problem are equally easy to solve.
The 2D case is much more interesting. Poisson's equation on the square with homogeneous Dirichlet conditions is:
$$
- \frac{d^2}{dx^2} u(x,y) - \frac{d^2}{dy^2}u(x,y) = f(x,y)\\
\qquad u(x,0) = u(x,1) =0\\
\qquad u(0,y) = u(1,y) = 0
$$
Suppose we discretize the equation above on a $n \times n$ uniform grid $(x_i, y_j)$, with $x_i = i h $, $y_j = jh$, with $h = \frac{1}{n+1}$.
We write $F_{i,j} = h^2 f(x_i, y_j)$, that is entry $(i,j)$ of the matrix $F \in \mathbb{R}^{n \times n}$ is the function $f(x,y)$ evaluated at $(x_i, y_j)$ times $h^2$.
We seek approximate value $U_{i,j} \approx u(x_i, y_j)$ through finite differences once again.
We can think of the operator $ u \rightarrow - \frac{d^2}{dx^2}u$ and boundary conditions $ u(x,0) = u(x,1) = 0 $ as being applied independently to the columns of $U$, thus they are discretized by the matrix $ U \rightarrow TU$.
Similarly, we can think of the operator $ u \rightarrow - \frac{d^2}{dy^2}u$ and boundary conditions $ u(0,y) = u(1,y) = 0 $ as being applied independently to the rows of $U$, thus they are discretized by the matrix $ U \rightarrow U T^T = UT$.
Hence, the operator
$ u \rightarrow - \frac{d^2}{dy^2}u - \frac{d^2}{dx^2}u$ and boundary conditions is discretized as $U \rightarrow TU + UT$, and we arrive at the matrix equation:
$$ UT + TU = F$$
Thus, the 2D Poisson finite-difference approximation may be computed by solving this matrix linear system for $U$.
Using the Kronecker formalism from the last HW, we could solve this in $O(n^6)$ operation: not great. Bartles-Stewart is much better, with $O(n^3)$ complexity, but we are not out of tricks just yet!
A wonderful property of the second-order finite difference matrix is that we can compute its eigenvalue decomposition analytically:
$$T = S^{-1} \Lambda S$$
where $S$ is the Discrete Sine Transform (DST) matrix defined componentwise by $S_{kj} = \frac{1}{\sqrt{2(n+1)}}\sin(\frac{kj\pi}{n+1})$, and eigenvalues are $\lambda_j = 4 \sin^2 ( \frac{j \pi}{2(n+1)})$. The DST is the imaginary part of the complex Discrete Fourier Matrix (DFT) $S = \text{imag}(F)$. Thanks to this connection, one can derive a $\mathcal{O}(n \log n)$ algorithm to compute $Sx$ based on the FFT, which can be obtained by scipy.fft.dst(x, type = 1, norm = "ortho"). The discrete Sine matrix has another nice property: it is both orthogonal and symmetric, thus $S = S^T = S^{-1}$ (when properly normalized, which is what the norm = "ortho" does).
Your task is to write a solver for Poisson's equation in $\mathcal{O}(n^2 \log n )$ using this property of the second-order finite difference matrix $T$ and the Bartels-Stewart algorithm.
(A programming note: if you want to compute $SX$ (where X is a matrix), you should use scipy.fft.dst(X, type = 1, norm = "ortho", axis =1). If you want to compute $XS^T$, you should compute scipy.fft.dst(X, type = 1, norm = "ortho", axis =0). You can also use the dst helper function below if you get tired of writing this long statement).
