```
import sympy
sympy.init_printing(use_latex='mathjax')
```

```
t0,t1,t2,t3,t4,t5 = sympy.symbols("t_0,t_1,t_2,t_3,t_4,t_5")
M = sympy.Matrix([[1, t0, t3], [0, t1, t4], [0, t2, t5]])
M
```

$$\left[\begin{matrix}1 & t_{0} & t_{3}\\0 & t_{1} & t_{4}\\0 & t_{2} & t_{5}\end{matrix}\right]$$

```
MI = M.inv()
MI
```

$$\left[\begin{matrix}1 & - \frac{t_{0}}{t_{1}} + \frac{t_{2} \left(- \frac{t_{0} t_{4}}{t_{1}} + t_{3}\right)}{t_{1} \left(t_{5} - \frac{t_{2} t_{4}}{t_{1}}\right)} & - \frac{- \frac{t_{0} t_{4}}{t_{1}} + t_{3}}{t_{5} - \frac{t_{2} t_{4}}{t_{1}}}\\0 & \frac{1}{t_{1}} + \frac{t_{2} t_{4}}{t_{1}^{2} \left(t_{5} - \frac{t_{2} t_{4}}{t_{1}}\right)} & - \frac{t_{4}}{t_{1} \left(t_{5} - \frac{t_{2} t_{4}}{t_{1}}\right)}\\0 & - \frac{t_{2}}{t_{1} \left(t_{5} - \frac{t_{2} t_{4}}{t_{1}}\right)} & \frac{1}{t_{5} - \frac{t_{2} t_{4}}{t_{1}}}\end{matrix}\right]$$

```
sympy.printing.print_tree(MI)
```

```
print repr(MI)
```

    Matrix([
    [1, -t_0/t_1 + t_2*(-t_0*t_4/t_1 + t_3)/(t_1*(t_5 - t_2*t_4/t_1)), -(-t_0*t_4/t_1 + t_3)/(t_5 - t_2*t_4/t_1)],
    [0,                  1/t_1 + t_2*t_4/(t_1**2*(t_5 - t_2*t_4/t_1)),            -t_4/(t_1*(t_5 - t_2*t_4/t_1))],
    [0,                                -t_2/(t_1*(t_5 - t_2*t_4/t_1)),                     1/(t_5 - t_2*t_4/t_1)]])

```
p,l,x, y = sympy.symbols("p,l,x,y")
e = sympy.Matrix([[1, x, y]]) * MI
print sympy.Eq(p, e.tolist()[0][1])
```

    p == -t_0/t_1 + x*(1/t_1 + t_2*t_4/(t_1**2*(t_5 - t_2*t_4/t_1))) - t_2*y/(t_1*(t_5 - t_2*t_4/t_1)) + t_2*(-t_0*t_4/t_1 + t_3)/(t_1*(t_5 - t_2*t_4/t_1))

```
print sympy.Eq(l, e.tolist()[0][2])
```

    l == y/(t_5 - t_2*t_4/t_1) - (-t_0*t_4/t_1 + t_3)/(t_5 - t_2*t_4/t_1) - t_4*x/(t_1*(t_5 - t_2*t_4/t_1))

```

```


# Diffusion in 2D

We now address diffusion in two space dimensions:

$$
\begin{equation}
\frac{\partial u}{\partial t}  = \dfc\left(
\frac{\partial^2 u}{\partial x^2} +
\frac{\partial^2 u}{\partial x^2}\right) + f(x,y),
\label{_auto1} \tag{1}
\end{equation}
$$

in a domain

$$
(x,y)\in (0,L_x)\times (0,L_y),\ t\in (0,T],
$$

with $u=0$ on the boundary and $u(x,y,0)=I(x,y)$ as initial condition.

## Discretization

For generality, it is natural to use a $\theta$-rule for the time
discretization. Standard, second-order accurate finite differences are
used for the spatial derivatives. We sample the PDE at a space-time
point $(i,j,n+\frac{1}{2})$ and apply the difference approximations:

$$
\lbrack D_t u\rbrack^{n+\frac{1}{2}} =
\theta \lbrack \dfc (D_xD_x u + D_yD_yu) + f\rbrack^{n+1} + \nonumber
$$

$$
\begin{equation}
\quad (1-\theta)\lbrack \dfc (D_xD_x u + D_yD_y u) + f\rbrack^{n}\thinspace .
\label{_auto2} \tag{2}
\end{equation}
$$

Written out,

$$
\frac{u^{n+1}_{i,j}-u^n_{i,j}}{\Delta t} =\nonumber
$$

$$
\qquad \theta (\dfc
(\frac{u^{n+1}_{i-1,j} - 2u^{n+1}_{i,j} + u^{n+1}_{i+1,j}}{\Delta x^2} +
\frac{u^{n+1}_{i,j-1} - 2u^{n+1}_{i,j} + u^{n+1}_{i,j+1}}{\Delta y^2}) +
f^{n+1}_{i,j})
+ \nonumber
$$

$$
\begin{equation}
\qquad (1-\theta)(\dfc
(\frac{u^{n}_{i-1,j} - 2u^{n}_{i,j} + u^{n}_{i+1,j}}{\Delta x^2} +
\frac{u^{n}_{i,j-1} - 2u^{n}_{i,j} + u^{n}_{i,j+1}}{\Delta y^2}) +
f^{n}_{i,j})
\label{_auto3} \tag{3}
\end{equation}
$$

We collect the unknowns on the left-hand side

$$
u^{n+1}_{i,j} -
\theta\left(
F_x
(u^{n+1}_{i-1,j} - 2u^{n+1}_{i,j} + u^{n+1}_{i+1,j}) +
F_y
(u^{n+1}_{i,j-1} - 2u^{n+1}_{i,j} + u^{n+1}_{i,j+1})\right)
= \nonumber
$$

$$
\qquad
(1-\theta)\left(
F_x
(u^{n}_{i-1,j} - 2u^{n}_{i,j} + u^{n}_{i+1,j}) +
F_y
(u^{n}_{i,j-1} - 2u^{n}_{i,j} + u^{n}_{i,j+1})\right) + \nonumber
$$

$$
\begin{equation}
\qquad \theta \Delta t f^{n+1}_{i,j} + (1-\theta) \Delta t f^{n}_{i,j}
+ u^n_{i,j},
\label{diffu:2D:theta_scheme2} \tag{4}
\end{equation}
$$

where

$$
F_x = \frac{\dfc\Delta t}{\Delta x^2},\quad F_y = \frac{\dfc\Delta t}{\Delta y^2},
$$

are the Fourier numbers in $x$ and $y$ direction, respectively.

3x2 2D mesh.

## Numbering of mesh points versus equations and unknowns

The equations ([4](#diffu:2D:theta_scheme2)) are coupled at the new
time level $n+1$. That is, we must solve a system of (linear) algebraic
equations, which we will write as $Ac=b$, where $A$ is the coefficient
matrix, $c$ is the vector of unknowns, and $b$ is the right-hand side.

Let us examine the equations in $Ac=b$ on a mesh with $N_x=3$ and
$N_y=2$ cells in the respective spatial directions.  The spatial mesh is depicted in
[Figure](#diffu:2D:fig:mesh3x2).  The equations at the boundary just
implement the boundary condition $u=0$:

$$
\begin{align*}
& u^{n+1}_{0,0}=
u^{n+1}_{1,0}=
u^{n+1}_{2,0}=
u^{n+1}_{3,0}=
u^{n+1}_{0,1}=\\
& u^{n+1}_{3,1}=
u^{n+1}_{0,2}=
u^{n+1}_{1,2}=
u^{n+1}_{2,2}=
u^{n+1}_{3,2}= 0\thinspace .
\end{align*}
$$

We are left with two interior points, with $i=1$, $j=1$ and $i=2$, $j=1$.
The corresponding equations are

$$
\begin{align*}
& u^{n+1}_{i,j} -
\theta\left(
F_x
(u^{n+1}_{i-1,j} - 2u^{n+1}_{i,j} + u^{n+1}_{i+1,j}) +
F_y
(u^{n+1}_{i,j-1} - 2u^{n+1}_{i,j} + u^{n+1}_{i,j+1})\right)
=  \\
&\qquad
(1-\theta)\left(
F_x
(u^{n}_{i-1,j} - 2u^{n}_{i,j} + u^{n}_{i+1,j}) +
F_y
(u^{n}_{i,j-1} - 2u^{n}_{i,j} + u^{n}_{i,j+1})\right) + \\
&\qquad \theta \Delta t f^{n+1}_{i,j} + (1-\theta) \Delta t f^{n}_{i,j}
+ u^n_{i,j},
\end{align*}
$$

There are in total 12 unknowns $u^{n+1}_{i,j}$ for $i=0,1,2,3$ and
$j=0,1,2$.  To solve the equations, we need to form a matrix system $Ac=b$.
In that system, the solution vector $c$ can only have one index. Thus,
we need a numbering of the unknowns with one
index, not two as used in the mesh. We introduce a mapping $m(i,j)$
from a mesh point with indices $(i,j)$ to the corresponding unknown
$p$ in the equation system:

$$
p = m(i,j) = j(N_x+1) + i\thinspace .
$$

When $i$ and $j$ run through their values, we see the following mapping
to $p$:

$$
\begin{align*}
&(0,0)\rightarrow 0,\
(0,1)\rightarrow 1,\
(0,2)\rightarrow 2,\
(0,3)\rightarrow 3,\\
&(1,0)\rightarrow 4,\
(1,1)\rightarrow 5,\
(1,2)\rightarrow 6,\
(1,3)\rightarrow 7,\\
&(2,0)\rightarrow 8,\
(2,1)\rightarrow 9,\
(2,2)\rightarrow 10,\
(2,3)\rightarrow 11\thinspace .
\end{align*}
$$

That is, we number the points along the $x$ axis, starting with $y=0$,
and then progress one "horizontal" mesh line at a time.
In [Figure](#diffu:2D:fig:mesh3x2) you can see that the $(i,j)$ and the
corresponding single index ($p$) are listed for each mesh point.

We could equally well have numbered the equations in other ways, e.g.,
let the $j$ index be the fastest varying index:
$p = m(i,j) = i(N_y+1) + j$.

Let us form the coefficient matrix $A$, or more precisely, insert a
matrix element (according Python's convention with zero as base
index) for each of the nonzero elements in $A$ (the indices
run through the values of $p$, i.e., $p=0,\ldots,11$):

$$
{\tiny
\left(\begin{array}{cccccccccccc}
(0,0) &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   & (1,1) &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   & (2,2) &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   & (3,3) &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   & (4,4) &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   & (5,1) &   0   &   0   & (5,4) & (5,5) & (5,6) &   0   &   0   & (5,9) &   0    &   0    \\
  0   &   0   & (6,2) &   0   &   0   & (6,5) & (6,6) & (6,7) &   0   &   0   & (6,10) &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   & (7,7) &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   & (8,8) &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   & (9,9) &   0    &   0    \\
   0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   & (10,10) &    0    \\
   0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0    & (11,11) \\
\end{array}\right)
}
$$

Here is a more compact visualization of the coefficient matrix where we
insert dots for zeros and bullets for non-zero elements:

$$
\footnotesize
\left(\begin{array}{cccccccccccc}
\bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \bullet & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \bullet & \cdot & \cdot \\
\cdot & \cdot & \bullet & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \bullet & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet \\
\end{array}\right)
$$

mathcal{I}_t is clearly seen that most of the elements are zero. This is a general
feature of coefficient matrices arising from discretizing PDEs by
finite difference methods. We say that the matrix is *sparse*.

Let $A_{p,q}$ be the value of element $(p,q)$ in the coefficient matrix $A$,
where $p$ and $q$ now correspond to the numbering of the unknowns in the
equation system.
We have $A_{p,q}=1$ for $p=q=0,1,2,3,4,7,8,9,10,11$, corresponding
to all the known boundary values. Let $p$ be $m(i,j)$, i.e.,
the single index corresponding to mesh point $(i,j)$. Then we have

$$
\begin{equation}
A_{m(i,j),m(i,j)} = A_{p,p} =
1 +
\theta (F_x + F_y),
\label{_auto4} \tag{5}
\end{equation}
$$

$$
\begin{equation}
A_{p, m(i-1,j)} = A_{p,p-1} = -\theta F_x,
\label{_auto5} \tag{6}
\end{equation}
$$

$$
\begin{equation}
A_{p, m(i+1,j)} = A_{p,p+1} = -\theta F_x,
\label{_auto6} \tag{7}
\end{equation}
$$

$$
\begin{equation}
A_{p, m(i,j-1)} = A_{p, p-(N_x+1)} = -\theta F_y,
\label{_auto7} \tag{8}
\end{equation}
$$

$$
\begin{equation}
A_{p, m(i,j+1)} = A_{p, p+(N_x+1)} = -\theta F_y,
\label{_auto8} \tag{9}
\end{equation}
$$

for the equations associated with the two interior mesh points.
At these interior points, the single index $p$ takes on the
specific values $p=5,6$, corresponding to the
values $(1,1)$ and $(1,2)$ of the pair $(i,j)$.

The above values for $A_{p,q}$ can be inserted in the matrix:

$$
{\tiny
\left(\begin{array}{cccccccccccc}
1 &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   & 1 &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   & 1 &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   & 1 &   0   &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   & 1 &   0   &   0   &   0   &   0   &   0   &   0    &   0    \\
  0   & -\theta F_y &   0   &   0   & -\theta F_x & 1+2\theta F_x & -\theta F_x &   0   &   0   & -\theta F_y &   0    &   0    \\
  0   &   0   & -\theta F_y &   0   &   0   & -\theta F_x & 1+2\theta F_x & -\theta F_x &   0   &   0   & -\theta F_y &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   & 1 &   0   &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   & 1 &   0   &   0    &   0    \\
  0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   &   0   & 1 &   0    &   0    \\
   0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   & 1 &    0    \\
   0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0   &    0    & 1 \\
\end{array}\right)
}
$$

The corresponding right-hand side vector in the equation system has
the entries $b_p$, where $p$ numbers the equations. We have

$$
b_0=b_1=b_2=b_3=b_4=b_7=b_8=b_9=b_{10}=b_{11}=0,
$$

for the boundary values. For the equations associated with the
interior points, we get for $p=5,6$, corresponding to $i=1,2$ and $j=1$:

$$
\begin{align*}
b_p &= u^{n}_{i,j} +
(1-\theta)\left(
F_x
(u^{n}_{i-1,j} - 2u^{n}_{i,j} + u^{n}_{i+1,j}) +
F_y
(u^{n}_{i,j-1} - 2u^{n}_{i,j} + u^{n}_{i,j+1})\right) + \\
&\qquad \theta \Delta t f^{n+1}_{i,j} + (1-\theta) \Delta t f^{n}_{i,j}\thinspace .
\end{align*}
$$

Recall that $p=m(i,j)=j(N_x+1)+j$ in this expression.

We can, as an alternative, leave the boundary mesh points out of the
matrix system. For a mesh with $N_x=3$ and $N_y=2$ there are only two
internal mesh points whose unknowns will enter the matrix system.
We must now number the unknowns at the interior points:

$$
p = (j-1)(N_x-1) + i,
$$

for $i=1,\ldots,N_x-1$, $j=1,\ldots,N_y-1$.

4x3 2D mesh.

We can continue with illustrating a bit larger mesh, $N_x=4$ and $N_y=3$,
see [Figure](#diffu:2D:fig:mesh4x3). The corresponding coefficient matrix
with dots for zeros and bullets for non-zeroes looks as follows (values at boundary points are included in the equation system):

$$
{\tiny
\left(\begin{array}{cccccccccccccccccccc}
\bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot \\
\cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet \\
\end{array}\right)
}
$$

**The coefficient matrix is banded.**

Besides being sparse, we observe that the coefficient matrix is *banded*:
it has five distinct bands. We have the diagonal $A_{i,i}$, the
subdiagonal $A_{i-1,j}$, the superdiagonal $A_{i,i+1}$, a lower
diagonal $A_{i,i-(Nx+1)}$, and an upper diagonal $A_{i,i+(Nx+1)}$.
The other matrix entries are known to be zero. With $N_x+1=N_y+1=N$,
only a fraction $5N^{-2}$ of the matrix entries are nonzero, so the
matrix is clearly very sparse for relevant $N$ values.
The more we can compute with the nonzeros only, the faster the solution
methods will potentially be.

## Algorithm for setting up the coefficient matrix

We looked at a specific mesh in the previous section, formulated
the equations, and saw what the corresponding coefficient matrix and
right-hand side are. Now our aim is to set up a general algorithm, for any
choice of $N_x$ and $N_y$, that produces the coefficient matrix and
the right-hand side vector.
We start with a zero matrix and vector, run through each mesh point,
and fill in the values depending on whether the mesh point is an interior
point or on the boundary.

 * for $i=0,\ldots,N_x$

  * for $j=0,\ldots, N_y$

    * $p=j(N_x+1)+i$

    * if point $(i,j)$ is on the boundary:

      * $A_{p,p}=1$, $b_p=0$

    * else:

      * fill $A_{p,m(i-1,j)}$, $A_{p,m(i+1,j)}$, $A_{p,m(i,j)}$, $A_{p,m(i,j-1)}$, $A_{p,m(i,j+1)}$, and $b_p$

To ease the test on whether $(i,j)$ is on the boundary or not, we can
split the loops a bit, starting with the boundary line $j=0$, then
treat the interior lines $1\leq j<N_y$, and finally treat the boundary
line $j=N_y$:

 * for $i=0,\ldots,N_x$

  * boundary $j=0$: $p=j(N_x+1)+i$, $A_{p,p}=1$

 * for $j=0,\ldots,N_y$

  * boundary $i=0$: $p=j(N_x+1)+i$, $A_{p,p}=1$

  * for $i=1,\ldots, N_x-1$

    * interior point $p=j(N_x+1)+i$

    * fill $A_{p,m(i-1,j)}$, $A_{p,m(i+1,j)}$, $A_{p,m(i,j)}$, $A_{p,m(i,j-1)}$, $A_{p,m(i,j+1)}$, and $b_p$

  * boundary $i=N_x$: $p=j(N_x+1)+i$, $A_{p,p}=1$

 * for $i=0,\ldots,N_x$

  * boundary $j=N_y$: $p=j(N_x+1)+i$, $A_{p,p}=1$

The right-hand side is set up as follows.

 * for $i=0,\ldots,N_x$

  * boundary $j=0$: $p=j(N_x+1)+i$, $b_p=0$

 * for $j=0,\ldots,N_y$

  * boundary $i=0$: $p=j(N_x+1)+i$, $b_p=0$

  * for $i=1,\ldots, N_x-1$

    * interior point $p=j(N_x+1)+i$

    * fill $b_p$

  * boundary $i=N_x$: $p=j(N_x+1)+i$, $b_p=0$

 * for $i=0,\ldots,N_x$

  * boundary $j=N_y$: $p=j(N_x+1)+i$, $b_p=0$

## Implementation with a dense coefficient matrix

The goal now is to map the algorithms in the previous section to
Python code. One should, for computational efficiency reasons, take
advantage of the fact that the coefficient matrix is sparse and/or
banded, i.e., take advantage of all the zeros. However, we first demonstrate
how to fill an $N\times N$ dense square matrix, where $N$ is the number
of unknowns, here $N=(N_x+1)(N_y+1)$. The dense matrix is much easier
to understand than the sparse matrix case.

```
import numpy as np

def solver_dense(
    I, a, f, Lx, Ly, Nx, Ny, dt, T, theta=0.5, user_action=None):
    """
    Solve u_t = a*(u_xx + u_yy) + f, u(x,y,0)=I(x,y), with u=0
    on the boundary, on [0,Lx]x[0,Ly]x[0,T], with time step dt,
    using the theta-scheme.
    """
    x = np.linspace(0, Lx, Nx+1)       # mesh points in x dir
    y = np.linspace(0, Ly, Ny+1)       # mesh points in y dir
    dx = x[1] - x[0]
    dy = y[1] - y[0]

    dt = float(dt)                    # avoid integer division
    Nt = int(round(T/float(dt)))
    t = np.linspace(0, Nt*dt, Nt+1)   # mesh points in time

    # Mesh Fourier numbers in each direction
    Fx = a*dt/dx**2
    Fy = a*dt/dy**2
```

The $u^{n+1}_{i,j}$ and $u^n_{i,j}$ mesh functions are represented
by their spatial values at the mesh points:

```
u   = np.zeros((Nx+1, Ny+1))      # unknown u at new time level
u_n = np.zeros((Nx+1, Ny+1))      # u at the previous time level
```

mathcal{I}_t is a good habit (for extensions) to
introduce index sets for all mesh points:

```
mathcal{I}_x = range(0, Nx+1)
Iy = range(0, Ny+1)
mathcal{I}_t = range(0, Nt+1)
```

The initial condition is easy to fill in:

```
# Load initial condition into u_n
for i in mathcal{I}_x:
    for j in Iy:
        u_n[i,j] = I(x[i], y[j])
```

The memory for the coefficient matrix and right-hand side vector
is allocated by

```
N = (Nx+1)*(Ny+1)  # no of unknowns
A = np.zeros((N, N))
b = np.zeros(N)
```

The filling of `A` goes like this:

```
m = lambda i, j: j*(Nx+1) + i

# Equations corresponding to j=0, i=0,1,... (u known)
j = 0
for i in mathcal{I}_x:
    p = m(i,j);  A[p, p] = 1

# Loop over all internal mesh points in y diretion
# and all mesh points in x direction
for j in Iy[1:-1]:
    i = 0;  p = m(i,j);  A[p, p] = 1   # Boundary
    for i in mathcal{I}_x[1:-1]:                 # Interior points
        p = m(i,j)
        A[p, m(i,j-1)] = - theta*Fy
        A[p, m(i-1,j)] = - theta*Fx
        A[p, p]        = 1 + 2*theta*(Fx+Fy)
        A[p, m(i+1,j)] = - theta*Fx
        A[p, m(i,j+1)] = - theta*Fy
    i = Nx;  p = m(i,j);  A[p, p] = 1  # Boundary
# Equations corresponding to j=Ny, i=0,1,... (u known)
j = Ny
for i in mathcal{I}_x:
    p = m(i,j);  A[p, p] = 1
```

Since `A` is independent of time, it can be filled once and for all before
the time loop. The right-hand side vector must be filled at each
time level inside the time loop:

```
import scipy.linalg

for n in mathcal{I}_t[0:-1]:
    # Compute b
    j = 0
    for i in mathcal{I}_x:
        p = m(i,j);  b[p] = 0           # Boundary
    for j in Iy[1:-1]:
        i = 0;  p = m(i,j);  b[p] = 0   # Boundary
        for i in mathcal{I}_x[1:-1]:              # Interior points
            p = m(i,j)
            b[p] = u_n[i,j] + \
              (1-theta)*(
              Fx*(u_n[i+1,j] - 2*u_n[i,j] + u_n[i-1,j]) +\
              Fy*(u_n[i,j+1] - 2*u_n[i,j] + u_n[i,j-1]))\
                + theta*dt*f(i*dx,j*dy,(n+1)*dt) + \
              (1-theta)*dt*f(i*dx,j*dy,n*dt)
        i = Nx;  p = m(i,j);  b[p] = 0  # Boundary
    j = Ny
    for i in mathcal{I}_x:
        p = m(i,j);  b[p] = 0           # Boundary

    # Solve matrix system A*c = b
    c = scipy.linalg.solve(A, b)

    # Fill u with vector c
    for i in mathcal{I}_x:
        for j in Iy:
            u[i,j] = c[m(i,j)]

    # Update u_n before next step
    u_n, u = u, u_n
```

We use `solve` from `scipy.linalg` and not from `numpy.linalg`. The difference
is stated below.

**`scipy.linalg` versus `numpy.linalg`.**

Quote from the [SciPy documentation](http://docs.scipy.org/doc/scipy/reference/tutorial/linalg.html):

`scipy.linalg` contains all the functions in `numpy.linalg`
plus some other more advanced ones not contained in `numpy.linalg`.

Another advantage of using `scipy.linalg` over numpy.linalg is that it is always compiled with BLAS/LAPACK support, while for NumPy this is optional. Therefore, the SciPy version might be faster depending on how NumPy was installed.

Therefore, unless you don't want to add SciPy as a dependency to your NumPy program, use `scipy.linalg` instead of `numpy.linalg`.

The code shown above is available in the `solver_dense` function
in the file [`diffu2D_u0.py`](${src_diffu}/diffu2D_u0.py), differing only
in the boundary conditions, which in the code can be an arbitrary function along
each side of the domain.

We do not bother to look at vectorized versions of filling `A` since
a dense matrix is just used of pedagogical reasons for the very first
implementation. Vectorization will be treated when `A` has a sparse
matrix representation, as in the section [Implementation with a sparse coefficient matrix](#diffu:2D:impl:sparse).

**How to debug the computation of $A$ and $b$.**

A good starting point for debugging the filling of $A$ and $b$ is
to choose a very coarse mesh, say $N_x=N_y=2$, where there is just
one internal mesh point, compute the equations by hand, and
print out `A` and `b` for comparison in the code. If wrong elements
in `A` or `b` occur, print out each assignment to elements in
`A` and `b` inside the loops and compare with what you expect.

To let the user store, analyze, or visualize the solution at each
time level, we include a callback function, named `user_action`,
to be called before the time loop and in each pass in that loop.
The function has the signature

```
user_action(u, x, xv, y, yv, t, n)
```

where `u` is a two-dimensional array holding the solution at time level `n`
and time `t[n]`. The $x$ and $y$ coordinates of the mesh points are given by
the arrays `x` and `y`, respectively. The arrays `xv` and `yv` are
vectorized representations of the mesh points such that vectorized
function evaluations can be invoked. The `xv` and `yv` arrays are
defined by

```
xv = x[:,np.newaxis]
yv = y[np.newaxis,:]
```

One can then evaluate, e.g., $f(x,y,t)$ at all internal mesh points at time
level `n` by first evaluating $f$ at all points,

```
f_a = f(xv, yv, t[n])
```

and then use slices to extract a view of the values at the internal
mesh points: `f_a[1:-1,1:-1]`.
The next section features an example on writing a `user_action`
callback function.

## Verification: exact numerical solution

A good test example to start with is one that preserves the solution
$u=0$, i.e., $f=0$ and $I(x,y)=0$. This trivial solution can uncover
some bugs.

The first real test example is based on having an exact solution of
the discrete equations. This solution is linear in time and quadratic
in space:

$$
u(x,y,t) = 5tx(L_x-x)y(y-L_y)\thinspace .
$$

Inserting this manufactured solution in the PDE shows that the
source term $f$ must be

$$
f(x,y,t) = 5x(L_x-x)y(y-L_y) + 10\dfc t (x(L_x-x)+ y(y-L_y))\thinspace .
$$

We can use the `user_action` function to compare the numerical solution
with the exact solution at each time level. A suitable helper function
for checking the solution goes like this:

        def quadratic(theta, Nx, Ny):

            def u_exact(x, y, t):
                return 5*t*x*(Lx-x)*y*(Ly-y)
            def I(x, y):
                return u_exact(x, y, 0)
            def f(x, y, t):
                return 5*x*(Lx-x)*y*(Ly-y) + 10*a*t*(y*(Ly-y)+x*(Lx-x))

            # Use rectangle to detect errors in switching i and j in scheme
            Lx = 0.75
            Ly = 1.5
            a = 3.5
            dt = 0.5
            T = 2

            def assert_no_error(u, x, xv, y, yv, t, n):
                """Assert zero error at all mesh points."""
                u_e = u_exact(xv, yv, t[n])
                diff = abs(u - u_e).max()
                tol = 1E-12
                msg = 'diff=%g, step %d, time=%g' % (diff, n, t[n])
                print msg
                assert diff < tol, msg

            solver_dense(
                I, a, f, Lx, Ly, Nx, Ny,
                dt, T, theta, user_action=assert_no_error)

A true test function for checking the quadratic solution for several
different meshes and $\theta$ values can take the form

```
def test_quadratic():
    # For each of the three schemes (theta = 1, 0.5, 0), a series of
    # meshes are tested (Nx > Ny and Nx < Ny)
    for theta in [1, 0.5, 0]:
        for Nx in range(2, 6, 2):
            for Ny in range(2, 6, 2):
                print 'testing for %dx%d mesh' % (Nx, Ny)
                quadratic(theta, Nx, Ny)
```

## Verification: convergence rates

For 2D verification with convergence rate computations, the expressions
and computations just build naturally on what we saw for 1D diffusion.
Truncation error analysis and other forms of error analysis point to a
numerical error formula like

$$
E = C_t\Delta t^p + C_x\Delta x^2 + C_y\Delta y^2,
$$

where $p$, $C_t$, $C_x$, and $C_y$ are constants. Often, the analysis of
a Crank-Nicolson
method can show that $p=2$, while the Forward and Backward Euler schemes
have $p=1$.

When checking the error formula empirically, we need to reduce it to
a form $E=Ch^r$ with a single discretization parameter $h$ and some
rate $r$ to be estimated. For the Backward Euler method,
where $p=1$, we can introduce a single discretization parameter
according to

$$
h = \Delta x^2 =  \Delta y^2,\quad h = K^{-1}\Delta t,
$$

where $K$ is a constant. The error formula then becomes

$$
E = C_t Kh + C_xh + C_yh = \tilde C h,\quad \tilde C = C_tK + C_x + C_y\thinspace .
$$

The simplest choice is obviously $K=1$. With the Forward Euler method, however,
stability requires $\Delta t = hK \leq h/(4\dfc)$, so $K\leq 1/(4\dfc)$.

For the Crank-Nicolson method, $p=2$, and we can simply choose

$$
h = \Delta x = \Delta y = \Delta t,
$$

since there is no restriction on $\Delta t$ in terms of $\Delta x$ and
$\Delta y$.

A frequently used error measure is the $\ell^2$ norm of the error mesh
point values. The section [wave:pde2:fd:MMS](#wave:pde2:fd:MMS) and the formula
([wave:pde2:fd:MMS:E:l2](#wave:pde2:fd:MMS:E:l2)) shows the error measure for a 1D
time-dependent problem. The extension to the current 2D problem
reads

$$
E = \left(\Delta t\Delta x\Delta y \sum_{n=0}^{N_t}
\sum_{i=0}^{N_x}\sum_{j=0}^{N_y}(\uex(x_i,y_j,t_n)
- u^n_{i,j})^2\right)^{\frac{1}{2}}\thinspace .
$$

One attractive manufactured solution is

$$
\uex = e^{-pt}\sin(k_xx)\sin(k_yy),\quad k_x=\frac{\pi}{L_x},
k_y=\frac{\pi}{L_y},
$$

where $p$ can be arbitrary. The required source term is

$$
f = (\dfc(k_x^2 + k_y^2) - p)\uex\thinspace .
$$

The function `convergence_rates` in
[`diffu2D_u0.py`](${src_diffu}/diffu2D_u0.py) implements a convergence
rate test. Two potential difficulties are important to be aware of:

1. The error formula is assumed to be
   correct when $h\rightarrow 0$, so for coarse meshes the estimated rate
   $r$ may be somewhat away from the expected value.
   Fine meshes may lead to prohibitively long execution times.

2. Choosing $p=\dfc (k_x^2 + k_y^2)$ in the manufactured solution above
   seems attractive ($f=0$), but leads to a slower approach to the
   asymptotic range where the error formula is valid (i.e., $r$
   fluctuates and needs finer meshes to stabilize).

## Implementation with a sparse coefficient matrix

We used a sparse matrix implementation in the section [diffu:pde1:impl:sparse](#diffu:pde1:impl:sparse)
for a 1D problem with a tridiagonal matrix. The present matrix, arising
from a 2D problem, has five diagonals, but we can use the same
sparse matrix data structure `scipy.sparse.diags`.

### Understanding the diagonals

Let us look closer at the diagonals in the example with a $4\times 3$ mesh
as depicted in [Figure](#diffu:2D:fig:mesh4x3) and its associated matrix
visualized by dots for zeros and bullets for nonzeros. From the example
mesh, we may generalize to an $N_x\times N_y$ mesh.

$$
{\tiny
\begin{array}{lcccccccccccccccccccc}
0 =m(0,0) & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
1 = m(1,0) & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
2 = m(2,0) & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
3 = m(3,0) & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
N_x=m(N_x,0) & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
N_x+1=m(0,1) & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
(N_x+1)+1=m(1,1) & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
(N_x+1)+2=m(2,1) & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
(N_x+1)+3=m(3,1) & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
(N_x+1)+N_x=m(N_x,1) & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
2(N_x+1)=m(0,2)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot \\
2(N_x+1)+1=m(1,2)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot \\
2(N_x+1)+2=m(2,2)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot \\
2(N_x+1)+3=m(3,2)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \bullet & \bullet & \bullet & \cdot & \cdot & \cdot & \bullet & \cdot \\
2(N_x+1)+N_x=m(N_x,2)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot & \cdot \\
N_y(N_x+1)=m(0,N_y)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot & \cdot \\
N_y(N_x+1)+1=m(1,N_y)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot & \cdot \\
N_y(N_x+1)+2=m(2,N_y)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot & \cdot \\
N_y(N_x+1)+3=m(3,N_y)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet & \cdot \\
N_y(N_x+1)+N_x=m(N_x,N_y)& \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \cdot & \bullet \\
\end{array}
}
$$

The main diagonal has $N=(N_x+1)(N_y+1)$ elements, while the sub- and
super-diagonals have $N-1$ elements. By looking at the matrix above,
we realize that the lower diagonal starts in row $N_x+1$ and goes to
row $N$, so its length is $N-(N_x+1)$. Similarly, the upper diagonal
starts at row 0 and lasts to row $N-(N_x+1)$, so it has the same length.
Based on this information, we declare the diagonals by

```
main   = np.zeros(N)            # diagonal
lower  = np.zeros(N-1)          # subdiagonal
upper  = np.zeros(N-1)          # superdiagonal
lower2 = np.zeros(N-(Nx+1))     # lower diagonal
upper2 = np.zeros(N-(Nx+1))     # upper diagonal
b      = np.zeros(N)            # right-hand side
```

### Filling the diagonals

We run through all mesh points and fill in elements on the various
diagonals. The line of mesh points corresponding to $j=0$ are all
on the boundary, and only the main diagonal gets a contribution:

```
m = lambda i, j: j*(Nx+1) + i
j = 0; main[m(0,j):m(Nx+1,j)] = 1  # j=0 boundary line
```

Then we run through all interior $j=\hbox{const}$ lines of mesh points.
The first and the last point on each line, $i=0$ and $i=N_x$, correspond
to boundary points:

```
for j in Iy[1:-1]:             # Interior mesh lines j=1,...,Ny-1
    i = 0;   main[m(i,j)] = 1
    i = Nx;  main[m(i,j)] = 1  # Boundary
```

For the interior mesh points $i=1,\ldots,N_x-1$ on a mesh line $y=\hbox{const}$
we can start with the main diagonal. The entries to be filled go from
$i=1$ to $i=N_x-1$ so the relevant slice in the `main` vector is
`m(1,j):m(Nx,j)`:

```
main[m(1,j):m(Nx,j)] = 1 + 2*theta*(Fx+Fy)
```

The `upper` array for the superdiagonal has its index 0 corresponding to
row 0 in the matrix, and the array entries
to be set go from $m(1,j)$ to $m(N_x-1,j)$:

```
upper[m(1,j):m(Nx,j)] = - theta*Fx
```

The subdiagonal (`lower` array), however, has its index 0
corresponding to row 1, so there is an offset of 1 in indices compared to
the matrix. The first nonzero occurs (interior point) at a mesh line $j=\hbox{const}$ corresponding to matrix row $m(1,j)$, and the corresponding array index
in `lower` is then $m(1,j)$. To fill the entries from $m(1,j)$ to $m(N_x-1,j)$
we set the following slice in `lower`:

```
lower_offset = 1
lower[m(1,j)-lower_offset:m(Nx,j)-lower_offset] = - theta*Fx
```

For the upper diagonal, its index 0 corresponds to matrix row 0, so there
is no offset and we can set the entries correspondingly to `upper`:

```
upper2[m(1,j):m(Nx,j)] = - theta*Fy
```

The `lower2` diagonal, however, has its first index 0 corresponding to row
$N_x+1$, so here we need to subtract the offset $N_x+1$:

```
lower2_offset = Nx+1
lower2[m(1,j)-lower2_offset:m(Nx,j)-lower2_offset] = - theta*Fy
```

We can now summarize the above code lines for setting the entries in
the sparse matrix representation of the coefficient matrix:

```
lower_offset = 1
lower2_offset = Nx+1
m = lambda i, j: j*(Nx+1) + i

j = 0; main[m(0,j):m(Nx+1,j)] = 1  # j=0 boundary line
for j in Iy[1:-1]:             # Interior mesh lines j=1,...,Ny-1
    i = 0;   main[m(i,j)] = 1  # Boundary
    i = Nx;  main[m(i,j)] = 1  # Boundary
    # Interior i points: i=1,...,N_x-1
    lower2[m(1,j)-lower2_offset:m(Nx,j)-lower2_offset] = - theta*Fy
    lower[m(1,j)-lower_offset:m(Nx,j)-lower_offset] = - theta*Fx
    main[m(1,j):m(Nx,j)] = 1 + 2*theta*(Fx+Fy)
    upper[m(1,j):m(Nx,j)] = - theta*Fx
    upper2[m(1,j):m(Nx,j)] = - theta*Fy
j = Ny; main[m(0,j):m(Nx+1,j)] = 1  # Boundary line
```

The next task is to create the sparse matrix from these diagonals:

```
import scipy.sparse

A = scipy.sparse.diags(
    diagonals=[main, lower, upper, lower2, upper2],
    offsets=[0, -lower_offset, lower_offset,
             -lower2_offset, lower2_offset],
    shape=(N, N), format='csr')
```

### Filling the right-hand side; scalar version

Setting the entries in the right-hand side is easier, since there are no
offsets in the array to take into account. The right-hand side is in fact similar to
the one previously shown, when we used a dense matrix representation
(the right-hand side vector is, of course, independent of what type of
representation we use for the coefficient matrix). The complete time
loop goes as follows.

```
import scipy.sparse.linalg

for n in mathcal{I}_t[0:-1]:
    # Compute b
    j = 0
    for i in mathcal{I}_x:
        p = m(i,j);  b[p] = 0                     # Boundary
    for j in Iy[1:-1]:
        i = 0;  p = m(i,j);  b[p] = 0             # Boundary
        for i in mathcal{I}_x[1:-1]:
            p = m(i,j)                            # Interior
            b[p] = u_n[i,j] + \
              (1-theta)*(
              Fx*(u_n[i+1,j] - 2*u_n[i,j] + u_n[i-1,j]) +\
              Fy*(u_n[i,j+1] - 2*u_n[i,j] + u_n[i,j-1]))\
                + theta*dt*f(i*dx,j*dy,(n+1)*dt) + \
              (1-theta)*dt*f(i*dx,j*dy,n*dt)
        i = Nx;  p = m(i,j);  b[p] = 0            # Boundary
    j = Ny
    for i in mathcal{I}_x:
        p = m(i,j);  b[p] = 0                     # Boundary

    # Solve matrix system A*c = b
    c = scipy.sparse.linalg.spsolve(A, b)

    # Fill u with vector c
    for i in mathcal{I}_x:
        for j in Iy:
            u[i,j] = c[m(i,j)]

    # Update u_n before next step
    u_n, u = u, u_n
```

### Filling the right-hand side; vectorized version

Since we use a sparse matrix and try to speed up the computations, we
should examine the loops and see if some can be easily removed by
vectorization. In the filling of $A$ we have already used vectorized
expressions at each $j=\hbox{const}$ line of mesh points. We can
very easily do the same in the code above and remove the need for
loops over the `i` index:

```
for n in mathcal{I}_t[0:-1]:
    # Compute b, vectorized version

    # Precompute f in array so we can make slices
    f_a_np1 = f(xv, yv, t[n+1])
    f_a_n   = f(xv, yv, t[n])

    j = 0; b[m(0,j):m(Nx+1,j)] = 0     # Boundary
    for j in Iy[1:-1]:
        i = 0;   p = m(i,j);  b[p] = 0 # Boundary
        i = Nx;  p = m(i,j);  b[p] = 0 # Boundary
        imin = mathcal{I}_x[1]
        imax = mathcal{I}_x[-1]  # for slice, max i index is mathcal{I}_x[-1]-1
        b[m(imin,j):m(imax,j)] = u_n[imin:imax,j] + \
              (1-theta)*(Fx*(
          u_n[imin+1:imax+1,j] -
        2*u_n[imin:imax,j] + \
          u_n[imin-1:imax-1,j]) +
                         Fy*(
          u_n[imin:imax,j+1] -
        2*u_n[imin:imax,j] +
          u_n[imin:imax,j-1])) + \
            theta*dt*f_a_np1[imin:imax,j] + \
          (1-theta)*dt*f_a_n[imin:imax,j]
    j = Ny;  b[m(0,j):m(Nx+1,j)] = 0 # Boundary

    # Solve matrix system A*c = b
    c = scipy.sparse.linalg.spsolve(A, b)

    # Fill u with vector c
    u[:,:] = c.reshape(Ny+1,Nx+1).T

    # Update u_n before next step
    u_n, u = u, u_n
```

The most tricky part of this code snippet is the loading of values from
the one-dimensional array `c`
into the two-dimensional array `u`. With our numbering of unknowns
from left to right along "horizontal" mesh lines, the correct
reordering of the one-dimensional array `c` as a two-dimensional array
requires first a reshaping to an `(Ny+1,Nx+1)` two-dimensional
array and then taking the transpose. The result is an `(Nx+1,Ny+1)`
array compatible with `u` both in size and appearance of the function values.

The `spsolve` function in `scipy.sparse.linalg` is an efficient
version of Gaussian elimination suited for matrices described by
diagonals.  The algorithm is known as *sparse Gaussian elimination*,
and `spsolve` calls up a well-tested C code called [SuperLU](http://crd-legacy.lbl.gov/~xiaoye/SuperLU/).

The complete code utilizing `spsolve`
is found in the `solver_sparse` function in the file
[`diffu2D_u0.py`](${src_diffu}/diffu2D_u0.py).

### Verification

We can easily extend the function `quadratic` from
the section [Verification: exact numerical solution](#diffu:2D:verify) to include a test of the
`solver_sparse` function as well.

```
def quadratic(theta, Nx, Ny):
    ...
    t, cpu = solver_sparse(
        I, a, f, Lx, Ly, Nx, Ny,
        dt, T, theta, user_action=assert_no_error)
```

## The Jacobi iterative method

So far we have created a matrix and right-hand side of a linear system
$Ac=b$ and solved the system for $c$ by calling an exact algorithm
based on Gaussian elimination. A much simpler implementation, which
requires no memory for the coefficient matrix $A$, arises if we solve
the system by *iterative* methods. These methods are only approximate,
and the core algorithm is repeated many times until the solution is
considered to be converged.

### Numerical scheme and linear system

To illustrate the idea of the Jacobi method, we simplify the numerical scheme to the
Backward Euler case, $\theta=1$, so there are fewer terms to write:

$$
u^{n+1}_{i,j} -
\left(
F_x
(u^{n+1}_{i-1,j} - 2u^{n+1}_{i,j} + u^{n+1}_{i+1,j}) +
F_y
(u^{n+1}_{i,j-1} - 2u^{n+1}_{i,j} + u^{n+1}_{i,j+1})\right)
= \nonumber
$$

$$
\begin{equation}  \qquad u^n_{i,j} + \Delta t f^{n+1}_{i,j}
\label{diffu:2D:BE_scheme} \tag{10}
\end{equation}
$$

The idea of the *Jacobi* iterative method is to introduce an iteration,
here with index $r$, where we in each iteration treat $u^{n+1}_{i,j}$
as unknown, but use values from the previous iteration for
the other unknowns $u^{n+1}_{i\pm 1,j\pm 1}$.

### mathcal{I}_terations

Let $u^{n+1,r}_{i,j}$
be the approximation to $u^{n+1}_{i,j}$ in iteration $r$, for all
relevant $i$ and $j$ indices. We first solve with respect to
$u^{n+1}_{i,j}$ to get the equation to solve:

$$
u^{n+1}_{i,j} = (1+2F_x +2F_y)^{-1}
\left(
F_x
(u^{n+1}_{i-1,j} + u^{n+1}_{i+1,j}) +
F_y
(u^{n+1}_{i,j-1} + u^{n+1}_{i,j+1})\right) +
\nonumber
$$

$$
\begin{equation}  \qquad
u^n_{i,j} + \Delta t f^{n+1}_{i,j}
\label{diffu:2D:Jacobi0} \tag{11}
\end{equation}
$$

The iteration is introduced by using iteration index $r$, for computed values,
on the right-hand side and $r+1$ (unknown in this iteration) on the left-hand
side:

$$
u^{n+1,r+1}_{i,j} = (1+2F_x +2F_y)^{-1}\left(
F_x
(u^{n+1,r}_{i-1,j} + u^{n+1,r}_{i+1,j}) +
F_y
(u^{n+1,r}_{i,j-1} + u^{n+1,r}_{i,j+1})\right)
\nonumber
$$

$$
\begin{equation}  \qquad
+ u^n_{i,j} + \Delta t f^{n+1}_{i,j}
\label{diffu:2D:Jacobi} \tag{12}
\end{equation}
$$

### Initial guess

We start the iteration with the computed values at the previous time level:

$$
\begin{equation}
u^{n+1,0}_{i,j} = u^{n}_{i,j},\quad i=0,\ldots,N_x,\ j=0,\ldots,N_y\thinspace .
\label{diffu:2D:iter:startvector} \tag{13}
\end{equation}
$$

### Relaxation

A common technique in iterative methods is to introduce a *relaxation*,
which means that the new approximation is a weighted mean of the
approximation as suggested by the algorithm and the previous approximation.
Naming the quantity on the left-hand side of ([12](#diffu:2D:Jacobi))
as $u^{n+1,*}_{i,j}$, a new approximation based on relaxation reads

$$
\begin{equation}
u^{n+1,r+1} = \omega u^{n+1,*}_{i,j} + (1-\omega) u^{n+1,r}_{i,j}\thinspace .
\label{diffu:2D:iter:relaxation} \tag{14}
\end{equation}
$$

Under-relaxation means $\omega < 1$, while over-relaxation has
$\omega > 1$.

### Stopping criteria

The iteration can be stopped when the change from one iteration to the
next is sufficiently small ($\leq \epsilon$), using either an infinity norm,

$$
\begin{equation}
\max_{i,j}\left\vert u^{n+1,r+1}_{i,j}-u^{n+1,r}_{i,j}
\right\vert \leq \epsilon,
\label{_auto9} \tag{15}
\end{equation}
$$

or an $L^2$ norm,

$$
\begin{equation}
\left(\Delta x\Delta y\sum_{i,j} (u^{n+1,r+1}_{i,j}-u^{n+1,r}_{i,j})^2
\right)^{\frac{1}{2}} \leq \epsilon\thinspace .
\label{_auto10} \tag{16}
\end{equation}
$$

Another widely used criterion measures how well the equations are solved
by looking at the residual (essentially $b-Ac^{r+1}$ if $c^{r+1}$ is
the approximation to the solution in iteration $r+1$).
The residual, defined in terms of the finite difference stencil, is

$$
R_{i,j} = u^{n+1,r+1}_{i,j} -
(F_x(u^{n+1,r+1}_{i-1,j} - 2u^{n+1,r+1}_{i,j} + u^{n+1,r+1}_{i+1,j}) +\nonumber
$$

$$
\quad\quad F_y(u^{n+1,r+1}_{i,j-1} - 2u^{n+1,r+1}_{i,j} + u^{n+1,r+1}_{i,j+1}))
- \nonumber
$$

$$
\begin{equation}
\qquad u^n_{i,j} -  \Delta t f^{n+1}_{i,j}
\label{diffu:2D:residual} \tag{17}
\end{equation}
$$

One can then iterate until the norm of the mesh function $R_{i,j}$
is less than some tolerance:

$$
\begin{equation}
\left(\Delta x\Delta y\sum_{i,j} R_{i,j}^2
\right)^{\frac{1}{2}} \leq \epsilon\thinspace .
\label{_auto11} \tag{18}
\end{equation}
$$

### Code-friendly notation

To make the mathematics as close as possible to what we will write in
a computer program, we may introduce some new notation: $u_{i,j}$ is a
short notation for $u^{n+1,r+1}_{i,j}$, $u^{-}_{i,j}$ is a short
notation for $u^{n+1,r}_{i,j}$, and $u^{(s)}_{i,j}$ denotes
$u^{n+1-s}_{i,j}$.  That is, $u_{i,j}$ is the unknown, $u^{-}_{i,j}$
is its most recently computed approximation, and $s$ counts time
levels backwards in time. The Jacobi method
([12](#diffu:2D:Jacobi))) takes the following form with the new
notation:

$$
u^{*}_{i,j} = (1+2F_x +2F_y)^{-1}((
F_x
(u^{-}_{i-1,j} + u^{-}_{i+1,j}) +
F_y
(u^{-}_{i,j-1} + u^{-}_{i,j+1})) +
\nonumber
$$

$$
\begin{equation}  \qquad
u^{(1)}_{i,j} + \Delta t f^{n+1}_{i,j})
\label{diffu:2D:Jacobi2} \tag{19}
\end{equation}
$$

### Generalization of the scheme

We can also quite easily introduce the $\theta$ rule for discretization in
time and write up the Jacobi iteration in that case as well:

$$
u^{*}_{i,j} = (1+ 2\theta(F_x +F_y))^{-1}(\theta(
F_x
(u^{-}_{i-1,j} + u^{-}_{i+1,j}) +
F_y
(u^{-}_{i,j-1} + u^{-}_{i,j+1})) +
\nonumber
$$

$$
\qquad
u^{(1)}_{i,j} + \theta \Delta t f^{n+1}_{i,j}
+ (1-\theta)\Delta t f^n_{i,j} + \nonumber
$$

$$
\begin{equation}
\qquad (1-\theta)(
F_x(u^{(1)}_{i-1,j}-2u^{(1)}_{i,j} + u^{(1)}_{i+1,j}) +
F_y(u^{(1)}_{i,j-1}-2u^{(1)}_{i,j} + u^{(1)}_{i,j+1})))\thinspace .
\label{diffu:2D:Jacobi3} \tag{20}
\end{equation}
$$

The final update of $u$ applies relaxation:

$$
u_{i,j} = \omega u^{*}_{i,j} + (1-\omega)u^{-}_{i,j}\thinspace .
$$

## Implementation of the Jacobi method

The Jacobi method needs no coefficient matrix and right-hand side
vector, but it needs an array for $u$ in the previous iteration.  We
call this array `u_`, using the notation at the end of the previous
section (at the same time level). The unknown itself is called `u`,
while `u_n` is the computed solution one time level back in time.
With a $\theta$ rule in time, the time loop can be coded like this:

```
for n in mathcal{I}_t[0:-1]:
    # Solve linear system by Jacobi iteration at time level n+1
    u_[:,:] = u_n  # Start value
    converged = False
    r = 0
    while not converged:
        if version == 'scalar':
            j = 0
            for i in mathcal{I}_x:
                u[i,j] = U_0y(t[n+1])           # Boundary
            for j in Iy[1:-1]:
                i = 0;   u[i,j] = U_0x(t[n+1])  # Boundary
                i = Nx;  u[i,j] = U_Lx(t[n+1])  # Boundary
		# Interior points
                for i in mathcal{I}_x[1:-1]:
                    u_new = 1.0/(1.0 + 2*theta*(Fx + Fy))*(theta*(
                        Fx*(u_[i+1,j] + u_[i-1,j]) +
                        Fy*(u_[i,j+1] + u_[i,j-1])) + \
                    u_n[i,j] + \
                    (1-theta)*(Fx*(
                    u_n[i+1,j] - 2*u_n[i,j] + u_n[i-1,j]) +
                      Fy*(
                    u_n[i,j+1] - 2*u_n[i,j] + u_n[i,j-1]))\
                      + theta*dt*f(i*dx,j*dy,(n+1)*dt) + \
                    (1-theta)*dt*f(i*dx,j*dy,n*dt))
                    u[i,j] = omega*u_new + (1-omega)*u_[i,j]
            j = Ny
            for i in mathcal{I}_x:
                u[i,j] = U_Ly(t[n+1])      # Boundary

        elif version == 'vectorized':
            j = 0;  u[:,j] = U_0y(t[n+1])  # Boundary
            i = 0;  u[i,:] = U_0x(t[n+1])  # Boundary
            i = Nx; u[i,:] = U_Lx(t[n+1])  # Boundary
            j = Ny; u[:,j] = U_Ly(t[n+1])  # Boundary
	    # Internal points
            f_a_np1 = f(xv, yv, t[n+1])
            f_a_n   = f(xv, yv, t[n])
            u_new = 1.0/(1.0 + 2*theta*(Fx + Fy))*(theta*(Fx*(
              u_[2:,1:-1] + u_[:-2,1:-1]) +
                Fy*(
              u_[1:-1,2:] + u_[1:-1,:-2])) +\
            u_n[1:-1,1:-1] + \
              (1-theta)*(Fx*(
              u_n[2:,1:-1] - 2*u_n[1:-1,1:-1] + u_n[:-2,1:-1]) +\
                Fy*(
              u_n[1:-1,2:] - 2*u_n[1:-1,1:-1] + u_n[1:-1,:-2]))\
              + theta*dt*f_a_np1[1:-1,1:-1] + \
              (1-theta)*dt*f_a_n[1:-1,1:-1])
            u[1:-1,1:-1] = omega*u_new + (1-omega)*u_[1:-1,1:-1]
        r += 1
        converged = np.abs(u-u_).max() < tol or r >= max_iter
        u_[:,:] = u

    # Update u_n before next step
    u_n, u = u, u_n
```

The vectorized version should be quite straightforward to understand
once one has an understanding of how a standard 2D finite stencil
is vectorized.

The first natural verification is to use the test problem
in the function `quadratic` from
the section [Verification: exact numerical solution](#diffu:2D:verify). This problem is known to have no
approximation error, but any iterative method will produce an
approximate solution with unknown error. For a tolerance $10^{-k}$
in the iterative method, we can, e.g., use a slightly larger
tolerance $10^{-(k-1)}$
for the difference between the exact and the computed solution.

```
def quadratic(theta, Nx, Ny):
    ...
    def assert_small_error(u, x, xv, y, yv, t, n):
        """Assert small error for iterative methods."""
        u_e = u_exact(xv, yv, t[n])
        diff = abs(u - u_e).max()
        tol = 1E-4
        msg = 'diff=%g, step %d, time=%g' % (diff, n, t[n])
        assert diff < tol, msg

    for version in 'scalar', 'vectorized':
        for theta in 1, 0.5:
            print 'testing Jacobi, %s version, theta=%g' % \
                  (version, theta)
            t, cpu = solver_Jacobi(
                I=I, a=a, f=f, Lx=Lx, Ly=Ly, Nx=Nx, Ny=Ny,
                dt=dt, T=T, theta=theta,
                U_0x=0, U_0y=0, U_Lx=0, U_Ly=0,
                user_action=assert_small_error,
                version=version, iteration='Jacobi',
                omega=1.0, max_iter=100, tol=1E-5)
```

Even for a very coarse $4\times 4$ mesh, the Jacobi method requires
26 iterations to reach a tolerance of $10^{-5}$,
which is quite many iterations, given that there are only 25 unknowns.

## Test problem: diffusion of a sine hill

mathcal{I}_t can be shown that

$$
\begin{equation}
\uex = Ae^{-\dfc\pi^2(L_x^{-2} + L_y^{-2})t}
\sin\left(\frac{\pi}{L_x}x\right)\sin\left(\frac{\pi}{L_y}y\right),
\label{diffu:2D:Jacobi:impl:hill:uex} \tag{21}
\end{equation}
$$

is a solution of the 2D homogeneous diffusion equation
$u_t = \dfc(u_{xx}+u_{yy})$ in
a rectangle $[0,L_x]\times [0,L_y]$, for any value of the amplitude $A$.
This solution vanishes at the boundaries,
and the initial condition is the product of two sines.
We may choose $A=1$ for simplicity.

mathcal{I}_t is difficult to know if our solver based on the Jacobi method works
properly since we are faced with two sources of errors: one from the
discretization, $E_\Delta$, and one from the iterative Jacobi method,
$E_i$. The total error in the computed $u$ can be represented as

$$
E_u = E_\Delta + E_i\thinspace .
$$

One error measure is to look at the maximum value, which is obtained for
the midpoint $x=L_x/2$ and $y=L_x/2$. This midpoint is represented in
the discrete `u` if $N_x$ and $N_y$ are even numbers. We can then
compute $E_u$ as $E_u = |\max \uex - \max u|$, when we know an exact
solution $\uex$ of the problem.

What about $E_\Delta$? If we use the maximum value as a measure of the
error, we have in fact analytical insight into the approximation error
in this particular problem. According to the section [diffu:2D:analysis](#diffu:2D:analysis), the exact solution
([21](#diffu:2D:Jacobi:impl:hill:uex)) of the PDE problem is also an
exact solution of the discrete equations, except that the damping
factor in time is different. More precisely,
([diffu:2D:analysis:BN:numexact](#diffu:2D:analysis:BN:numexact)) and
([diffu:2D:analysis:CN:numexact](#diffu:2D:analysis:CN:numexact)) are solutions of the discrete
problem for $\theta=1$ (Backward Euler) and $\theta=\frac{1}{2}$
(Crank-Nicolson), respectively.  The factors raised to the power $n$
is the numerical amplitude, and the errors in these factors become

$$
\begin{align*}
E_\Delta &= e^{-\dfc k^2t} - \left(
\frac{1 - 2(F_x\sin^2 p_x + F_x\sin^2p_y)}{1 + 2(F_x\sin^2 p_x + F_x\sin^2p_y)}
\right)^n,\quad \theta=\frac{1}{2},\\
E_\Delta &= e^{-\dfc k^2t} -
(1 + 4F_x\sin^2 p_x + 4F_y\sin^2 p_y)^{-n},\quad\theta=1\thinspace .
\end{align*}
$$

We are now in a position to compute $E_i$ numerically. That is, we can
compute the error due to iterative solution of the linear system and
see if it corresponds to the convergence tolerance used in the method.
Note that the convergence is based on measuring the difference in
two consecutive approximations, which is not exactly the error
due to the iteration, but it is a kind of measure, and it should
have about the same size as $E_i$.

The function `demo_classic_iterative` in [`diffu2D_u0.py`](${src_diffu}/diffu2D_u0.py) implements the idea above (also for the
methods in the section [The Gauss-Seidel and SOR methods](#diffu:2D:SOR)). The value of $E_i$ is in
particular printed at each time level. By changing the tolerance in
the convergence criterion of the Jacobi method, we can see that $E_i$
is of the same order of magnitude as the prescribed tolerance in the
Jacobi method.  For example: $E_\Delta\sim 10^{-2}$ with $N_x=N_y=10$
and $\theta=\frac{1}{2}$, as long as $\max u$ has some significant size
($\max u > 0.02$). An appropriate value of the tolerance is then
$10^{-3}$, such that the error in the Jacobi method does not become
bigger than the discretization error. In that case, $E_i$ is around
$5\cdot 10^{-3}$.  The corresponding number of Jacobi iterations (with
$\omega=1$) varies from 31 to 12 during the time simulation (for $\max
u > 0.02$). Changing the tolerance to $10^{-5}$ causes many more
iterations (61 to 42) without giving any contribution to the overall
accuracy, because the total error is dominated by $E_\Delta$.

Also, with an $N_x=N_y=20$, the spatial accuracy increases and many
more iterations are needed (143 to 45), but the dominating error is
from the time discretization. However, with such a finer spatial mesh,
a higher tolerance in the convergence criterion $10^{-4}$ is needed to
keep $E_i\sim 10^{-3}$.  More experiments show the disadvantage of the
very simple Jacobi iteration method: the number of iterations
increases with the number of unknowns, keeping the tolerance fixed,
but the tolerance should also be lowered to avoid the iteration error
to dominate the total error. A small adjustment of the Jacobi method,
as described in the section [The Gauss-Seidel and SOR methods](#diffu:2D:SOR), provides a better method.

## The relaxed Jacobi method and its relation to the Forward Euler method

We shall now show that solving the Poisson equation $-\dfc\nabla^2 u =
f$ by the Jacobi iterative method is in fact equivalent to using a
Forward Euler scheme on $u_t = \dfc\nabla^2 u + f$ and letting
$t\rightarrow\infty$.

A Forward Euler discretization of the 2D diffusion equation,

$$
\lbrack D_t^+ u = \dfc (D_xD_x u + D_yD_y u) + f\rbrack^n_{i,j},
$$

can be written out as

$$
u_{i,j}^{n+1} = u_{i,j}^n + \frac{\Delta t}{\dfc h^2}
\left( u_{i-1,j}^n + u_{i+1,j}^n + u_{i,j-1}^n + u_{i,j+1}^n - 4u_{i,j}^n
+ h^2f_{i,j}\right),
$$

where $h=\Delta x = \Delta y$ has been introduced for simplicity. The
scheme can be reordered as

$$
u_{i,j}^{n+1} = \left(1 - \omega\right) u_{i,j}^n
+ \frac{1}{4}\omega
\left( u_{i-1,j}^n + u_{i+1,j}^n + u_{i,j-1}^n + u_{i,j+1}^n - 4u_{i,j}^n
+ h^2f_{i,j}\right),
$$

with

$$
\omega = 4\frac{\Delta t}{\dfc h^2},
$$

but this latter form is nothing but the relaxed Jacobi method applied to

$$
[D_xD_x u + D_yD_y u = -f]^n_{i,j}\thinspace .
$$

From the equivalence above we know a couple of things about the Jacobi
method for solving $-\nabla^2 u = f$:

1. The method is unstable if $\omega > 1$ (since the Forward Euler method
   is then unstable).

2. The convergence is really slow as the iteration index increases (coming from the fact that
   the Forward Euler scheme requires many small time steps to reach the stationary solution).

These observations are quite disappointing: if we already have a
time-dependent diffusion problem and want to take larger time steps by
an implicit time discretization method, we will with the Jacobi method
end up with something close to a slow Forward Euler simulation of the
original problem at each time level.  Nevertheless, the are two
reasons for why the Jacobi method remains a fundamental building block
for solving linear systems arising from PDEs: 1) a couple of
iterations remove large parts of the error and this is effectively
used in the very efficient class of multigrid methods; and 2) the idea
of the Jacobi method can be developed into more efficient methods,
especially the SOR method, which is treated next.

## The Gauss-Seidel and SOR methods

If we update the mesh points according to the Jacobi method
([11](#diffu:2D:Jacobi0)) for a Backward Euler discretization with a
loop over $i=1,\ldots,N_x-1$ and $j=1,\ldots,N_y-1$, we realize that
when $u^{n+1,r+1}_{i,j}$ is computed, $u^{n+1,r+1}_{i-1,j}$ and
$u^{n+1,r+1}_{i,j-1}$ are already computed, so these new values can be
used rather than $u^{n+1,r}_{i-1,j}$ and $u^{n+1,r}_{i,j-1}$
(respectively) in the formula for $u^{n+1,r+1}_{i,j}$.  This idea
gives rise to the *Gauss-Seidel* iteration method, which
mathematically is just a small adjustment of ([11](#diffu:2D:Jacobi0)):

$$
u^{n+1,r+1}_{i,j} = (1+2F_x +2F_y)^{-1}((\nonumber
$$

$$
\begin{equation}  \qquad
F_x
(u^{n+1,r+1}_{i-1,j} + u^{n+1,r}_{i+1,j}) +
F_y
(u^{n+1,r+1}_{i,j-1} + u^{n+1,r}_{i,j+1})) +
u^n_{i,j} + \Delta t f^{n+1}_{i,j})\thinspace .
\label{diffu:2D:SOR:eq} \tag{22}
\end{equation}
$$

Observe that the way we access the mesh points in the formula
([22](#diffu:2D:SOR:eq)) is important: points with $i-1$ must be
computed before points with $i$, and points with $j-1$ must be
computed before points with $j$. Any sequence of mesh points can be
used in the Gauss-Seidel method, but the particular math formula must
distinguish between already visited points in the current iteration
and the points not yet visited.

The idea of relaxation ([14](#diffu:2D:iter:relaxation)) can equally
well be applied to the Gauss-Seidel method. Actually, the Gauss-Seidel
method with an arbitrary $0<\omega\leq 2$ has its own name: the
*Successive Over-Relaxation* method, abbreviated as SOR.

The SOR method for a $\theta$ rule discretization, with the
shortened $u$ and $u^{-}$ notation, can be written

$$
u^{*}_{i,j} = (1+ 2\theta(F_x +F_y))^{-1}(\theta(
F_x
(u_{i-1,j} + u^{-}_{i+1,j}) +
F_y
(u_{i,j-1} + u^{-}_{i,j+1})) +
\nonumber
$$

$$
\qquad
u^{(1)}_{i,j} + \theta \Delta t f^{n+1}_{i,j}
+ (1-\theta)\Delta t f^n_{i,j} + \nonumber
$$

$$
\begin{equation}
\qquad (1-\theta)(
F_x(u^{(1)}_{i-1,j}-2u^{(1)}_{i,j} + u^{(1)}_{i+1,j}) +
F_y(u^{(1)}_{i,j-1}-2u^{(1)}_{i,j} + u^{(1)}_{i,j+1}))),
\label{diffu:2D:SOR3} \tag{23}
\end{equation}
$$

$$
\begin{equation}
u_{i,j} = \omega u^{*}_{i,j} + (1-\omega)u^{-}_{i,j}
\label{_auto12} \tag{24}
\end{equation}
$$

The sequence of mesh points in ([23](#diffu:2D:SOR3)) is
$i=1,\ldots,N_x-1$, $j=1,\ldots,N_y-1$ (but whether $i$ runs faster
or slower than $j$ does not matter).

## Scalar implementation of the SOR method

Since the Jacobi and Gauss-Seidel methods with relaxation
are so similar, we can easily make a common code for the two:

```
for n in mathcal{I}_t[0:-1]:
    # Solve linear system by Jacobi/SOR iteration at time level n+1
    u_[:,:] = u_n  # Start value
    converged = False
    r = 0
    while not converged:
        if version == 'scalar':
            if iteration == 'Jacobi':
                u__ = u_
            elif iteration == 'SOR':
                u__ = u
            j = 0
            for i in mathcal{I}_x:
                u[i,j] = U_0y(t[n+1])  # Boundary
            for j in Iy[1:-1]:
                i = 0;   u[i,j] = U_0x(t[n+1])  # Boundary
                i = Nx;  u[i,j] = U_Lx(t[n+1])  # Boundary
                for i in mathcal{I}_x[1:-1]:
                    u_new = 1.0/(1.0 + 2*theta*(Fx + Fy))*(theta*(
                        Fx*(u_[i+1,j] + u__[i-1,j]) +
                        Fy*(u_[i,j+1] + u__[i,j-1])) + \
                    u_n[i,j] + (1-theta)*(
                      Fx*(
                    u_n[i+1,j] - 2*u_n[i,j] + u_n[i-1,j]) +
                      Fy*(
                    u_n[i,j+1] - 2*u_n[i,j] + u_n[i,j-1]))\
                      + theta*dt*f(i*dx,j*dy,(n+1)*dt) + \
                    (1-theta)*dt*f(i*dx,j*dy,n*dt))
                    u[i,j] = omega*u_new + (1-omega)*u_[i,j]
                j = Ny
                for i in mathcal{I}_x:
                    u[i,j] = U_Ly(t[n+1])  # boundary
        r += 1
        converged = np.abs(u-u_).max() < tol or r >= max_iter
        u_[:,:] = u

    u_n, u = u, u_n  # Get ready for next iteration
```

The idea here is to introduce `u__` to be used for already computed
values (`u`) in the Gauss-Seidel/SOR version of the implementation, or
just values from the previous iteration (`u_`) in case of the Jacobi method.

## Vectorized implementation of the SOR method

Vectorizing the Gauss-Seidel iteration step turns out to be non-trivial.
The problem is that vectorized operations typically imply
operations on arrays where the sequence in which we visit the elements does
not matter. In particular, this principle makes vectorized code trivial to
parallelize. However, in the Gauss-Seidel algorithm, the sequence in which we
visit the elements in the arrays does matter, and it is well known that
the basic method as explained above cannot be parallelized.
Therefore, also vectorization will require new thinking.

The strategy for vectorizing (and parallelizing) the Gauss-Seidel
method is to use a special numbering of the mesh points called
red-black numbering: every other point is red or black as in a
checkerboard pattern. This numbering requires $N_x$ and $N_y$ to
be even numbers.
Here is an example of a $6\times 6$ mesh:

        r b r b r b r
        b r b r b r b
        r b r b r b r
        b r b r b r b
        r b r b r b r
        b r b r b r b
        r b r b r b r

The idea now is to first update all the red points. Each formula for
updating a red point involves only the black neighbors. Thereafter, we
update all the black points, and at each black point, only the
recently computed red points are involved.

The scalar implementation of the red-black numbered Gauss-Seidel
method is really compact, since we can update values directly in
`u` (this guarantees that we use the most recently computed
values). Here is the relevant code for the Backward Euler
scheme in time and without a source term:

```
# Update internal points
for sweep in 'red', 'black':
    for j in range(1, Ny, 1):
        if sweep == 'red':
            start = 1 if j % 2 == 1 else 2
        elif sweep == 'black':
            start = 2 if j % 2 == 1 else 1
        for i in range(start, Nx, 2):
	    u[i,j] = 1.0/(1.0 + 2*(Fx + Fy))*(
                     Fx*(u[i+1,j] + u[i-1,j]) +
                     Fy*(u[i,j+1] + u[i,j-1]) + u_n[i,j])
```

The vectorized version must be based on slices. Looking at a typical
red-black pattern, e.g.,

        r b r b r b r
        b r b r b r b
        r b r b r b r
        b r b r b r b
        r b r b r b r
        b r b r b r b
        r b r b r b r

we want to update the internal points (marking boundary points with
`x`):

        x x x x x x x
        x r b r b r x
        x b r b r b x
        x r b r b r x
        x b r b r b x
        x r b r b r x
        x x x x x x x

mathcal{I}_t is impossible to make one slice that picks out all the internal
red points. Instead, we need two slices. The first involves points
marked with `R`:

        x x x x x x x
        x R b R b R x
        x b r b r b x
        x R b R b R x
        x b r b r b x
        x R b R b R x
        x x x x x x x

This slice is specified as `1::2` for `i` and `1::2` for `j`, or with
`slice` objects:

```
i = slice(1, None, 2);  j = slice(1, None, 2)
```

The second slice involves the red points with `R`:

        x x x x x x x
        x r b r b r x
        x b R b R b x
        x r b r b r x
        x b R b R b x
        x r b r b r x
        x x x x x x x

The slices are

```
i = slice(2, None, 2);  j = slice(2, None, 2)
```

For the black points, the first slice involves the `B` points:

        x x x x x x x
        x r B r B r x
        x b r b r b x
        x r B r B r x
        x b r b r b x
        x r B r B r x
        x x x x x x x

with slice objects

```
i = slice(2, None, 2);  j = slice(1, None, 2)
```

The second set of black points is shown here:

        x x x x x x x
        x r b r b r x
        x B r B r B x
        x r b r b r x
        x B r B r B x
        x r b r b r x
        x x x x x x x

with slice objects

```
i = slice(1, None, 2);  j = slice(2, None, 2)
```

That is, we need four sets of slices. The simplest way of implementing
the algorithm is to make a function with variables for the slices
representing $i$, $i-1$, $i+1$, $j$, $j-1$, and $j+1$, here called
`ic` ("i center"), `im1` ("i minus 1", `ip1` ("i plus 1"), `jc`, `jm1`,
and `jp1`, respectively.

```
def update(u_, u_n, ic, im1, ip1, jc, jm1, jp1):
    return \
       1.0/(1.0 + 2*theta*(Fx + Fy))*(theta*(
           Fx*(u_[ip1,jc] + u_[im1,jc]) +
           Fy*(u_[ic,jp1] + u_[ic,jm1])) +\
       u_n[ic,jc] + (1-theta)*(
         Fx*(u_n[ip1,jc] - 2*u_n[ic,jc] + u_n[im1,jc]) +\
         Fy*(u_n[ic,jp1] - 2*u_n[ic,jc] + u_n[ic,jm1]))+\
         theta*dt*f_a_np1[ic,jc] + \
         (1-theta)*dt*f_a_n[ic,jc])
```

The formula returned from `update` is to be compared with
([23](#diffu:2D:SOR3)).

The relaxed Jacobi iteration can be implemented by

```
ic  = jc  = slice(1,-1)
im1 = jm1 = slice(0,-2)
ip1 = jp1 = slice(2,None)
u_new[ic,jc] = update(
    u_, u_n, ic, im1, ip1, jc, jm1, jp1)
u[ic,jc] = omega*u_new[ic,jc] + (1-omega)*u_[ic,jc]
```

The Gauss-Seidel (or SOR) updates need four different steps.
The `ic` and `jc` slices are specified above. For each of these,
we must specify the corresponding `im1`, `ip1`, `jm1`, and `jp1`
slices. The code below contains the details.

```
# Red points
ic  = slice(1,-1,2)
im1 = slice(0,-2,2)
ip1 = slice(2,None,2)
jc  = slice(1,-1,2)
jm1 = slice(0,-2,2)
jp1 = slice(2,None,2)
u_new[ic,jc] = update(
    u_new, u_n, ic, im1, ip1, jc, jm1, jp1)

ic  = slice(2,-1,2)
im1 = slice(1,-2,2)
ip1 = slice(3,None,2)
jc  = slice(2,-1,2)
jm1 = slice(1,-2,2)
jp1 = slice(3,None,2)
u_new[ic,jc] = update(
    u_new, u_n, ic, im1, ip1, jc, jm1, jp1)

# Black points
ic  = slice(2,-1,2)
im1 = slice(1,-2,2)
ip1 = slice(3,None,2)
jc  = slice(1,-1,2)
jm1 = slice(0,-2,2)
jp1 = slice(2,None,2)
u_new[ic,jc] = update(
    u_new, u_n, ic, im1, ip1, jc, jm1, jp1)

ic  = slice(1,-1,2)
im1 = slice(0,-2,2)
ip1 = slice(2,None,2)
jc  = slice(2,-1,2)
jm1 = slice(1,-2,2)
jp1 = slice(3,None,2)
u_new[ic,jc] = update(
    u_new, u_n, ic, im1, ip1, jc, jm1, jp1)

# Relax
c = slice(1,-1)
u[c,c] = omega*u_new[c,c] + (1-omega)*u_[c,c]
```

The function `solver_classic_iterative` in
[`diffu2D_u0.py`](${src_diffu}/diffu2D_u0.py)
contains a unified implementation of the relaxed Jacobi and SOR
methods in scalar and vectorized versions using the techniques
explained above.

## Direct versus iterative methods

### Direct methods

There are two classes of methods for solving linear systems: direct methods
and iterative methods. Direct methods are based on variants of the
Gaussian elimination procedure and will produce an exact solution (in
exact arithmetics) in an a priori known number of steps.
mathcal{I}_terative methods, on the other hand, produce an approximate solution,
and the amount of work for reaching a given accuracy is usually not
known.

The most common direct method today is to use the *LU factorization*
procedure to factor the coefficient matrix $A$ as the product of a
lower-triangular matrix $L$ (with unit diagonal terms)
and an upper-triangular matrix $U$:
$A=LU$. As soon as we have $L$ and $U$, a system of equations $LUc=b$
is easy to solve because of the triangular nature of $L$ and $U$. We
first solve $Ly=b$ for $y$ (forward substitution), and
thereafter we find $c$ from solving $Uc=y$ (backward
substitution).  When $A$ is a dense $N\times N$ matrix, the LU
factorization costs $\frac{1}{3}N^3$ arithmetic operations, while the
forward and backward substitution steps each require of
the order $N^2$ arithmetic operations.
That is, factorization dominates the costs, while the substitution steps are
cheap.

Symmetric, positive definite coefficient matrices often arise when
discretizing PDEs. In this case, the LU factorization becomes $A=LL^T$,
and the associated algorithm is known as *Cholesky factorization*.
Most linear algebra software offers highly optimized implementations of
LU and Cholesky factorization as well as forward and backward
substitution (`scipy.linalg` is the relevant Python package).

Finite difference discretizations lead to sparse coefficient matrices.
An extreme case arose in the section [diffu:pde1:BE](#diffu:pde1:BE) where $A$ was
tridiagonal. For a tridiagonal matrix, the amount of arithmetic
operations in the LU and Cholesky factorization algorithms is just of
the order $N$, not $N^3$. Tridiagonal matrices are special cases of
*banded matrices*, where the matrices contain just a set of diagonal
bands.  Finite difference methods on regularly numbered rectangular
and box-shaped meshes give rise to such banded matrices, with 5 bands
in 2D and 7 in 3D for diffusion problems.  Gaussian elimination only
needs to work within the bands, leading to much more efficient
algorithms.

If $A_{i,j}=0$ for $j> i+p$ and $j< i-p$, $p$ is the *frac{1}{2}-bandwidth*
of the matrix. We have in our 2D problem $p=N_x+2$, while in 3D,
$p=(N_x+1)(N_y+1)+2$. The cost of Gaussian elimination is then
$\Oof{Np^2}$, so with $p\ll N$, we see that banded matrices are much
more efficient to compute with. By reordering the unknowns in clever
ways, one can reduce the work of Gaussian elimination
further. Fortunately, the Python programmer has access to such
algorithms through the `scipy.sparse.linalg` package.

Although a direct method is an exact algorithm, rounding errors
may in practice accumulate and pollute the solution. The effect
grows with the size of the linear system, so both for accuracy and
efficiency, iterative methods are better suited than direct methods
for solving really large linear systems.

### mathcal{I}_terative methods

The Jacobi and SOR iterative methods belong to a class of iterative methods
where the idea is to solve $Au=b$ by splitting A into two parts, $A=M-N$,
such that solving systems $Mu=c$ is easy and efficient. With the splitting,
we get a system

$$
Mu = Nu + b,
$$

which suggests an iterative method

$$
Mu^{r+1} = Nu^{r} + b,\quad r=0,1,2,\ldots,
$$

where $u^{r+1}$ is a new approximation to $u$ in the $r+1$-th iteration. To
initiate the iteration, we need a start vector $u^0$.

The Jacobi and SOR methods are based on splitting $A$ into a lower
tridiagonal part $L$, the diagonal $D$, and an upper tridiagonal part $U$,
such that $A=L+D+U$. The Jacobi method corresponds to $M=D$ and $N=-L-U$.
The Gauss-Seidel method employs $M=L+D$ and $N=-U$, while the SOR
method corresponds to

$$
M= \frac{1}{\omega}D + L,\quad N = \frac{1-\omega}{\omega}D - U\thinspace .
$$

The relaxed Jacobi method has similar expressions:

$$
M = \frac{1}{\omega}D,\quad N = \frac{1-\omega}{\omega}D - L - U\thinspace .
$$

With the matrix forms of the Jacobi and SOR methods as written above,
we could in an implementation alternatively fill the matrix $A$ with
entries and call general implementations of the Jacobi or SOR methods
that work on a system $Au=b$. However, this is almost never done since
forming the matrix $A$ requires quite some code and storing $A$ in the
computer's memory is unnecessary. mathcal{I}_t is much easier to just apply the
Jacobi and SOR ideas to the finite difference stencils directly in an
implementation, as we have shown in detail.

Nevertheless, the matrix formulation of the Jacobi and SOR methods have
been important for analyzing their convergence behavior. One can show that
the error $u^r-u$ fulfills $u^r-u = G^r(u^0-u)$, where $G=M^{-1}N$
and $G^k$ is a matrix exponential. For the method to converge,
$\lim_{r\rightarrow\infty}||G^r||=0$ is a necessary and sufficient condition.
This implies that the *spectral radius* of $G$ must be less than one.
Since $G$ is directly related to the finite difference scheme for the
underlying PDE problem, one can in principle compute the spectral radius.
For a given PDE problem, however, this is not a practical strategy, since it
is very difficult to develop useful formulas.
Analysis of model problems, usually
related to the Poisson equation,
reveals some trends of interest: the convergence rate of the Jacobi
method goes like $h^2$, while that of SOR with an optimal $\omega$ goes
like $h$, where $h$ is the spatial spacing: $h=\Delta x=\Delta y$.
That is, the efficiency of the Jacobi method quickly deteriorates with
the increasing mesh resolution, and SOR is much to be preferred
(even if the optimal
$\omega$ remains an open question). We refer to
Chapter 4 of [[Saad_2003]](#Saad_2003) for more information on the convergence theory.
One important result is that if $A$ is symmetric and positive definite,
then SOR will converge for any $0<\omega <2$.

The optimal $\omega$ parameter can be theoretically established for
a Poisson problem as

$$
\begin{equation}
\omega_{o} = \frac{2}{1 + \sqrt{1-\varrho^2}},\quad \varrho = \frac{\cos(\pi/N_x) + (\Delta x/\Delta y)^2\cos(\pi/N_y)}{1 + (\Delta x/\Delta y)^2}\thinspace .
\label{_auto13} \tag{25}
\end{equation}
$$

This formula can be used as a guide also in other problems.

The Jacobi and the SOR methods have their great advantage of being
trivial to implement, so they are obviously popular of this reason.
However, the slow convergence of these methods limits the popularity to fairly
small linear systems (i.e., coarse meshes). As soon as the matrix size
grows, one is better off with more sophisticated iterative methods
like the preconditioned Conjugate gradient method, which we now turn to.

Finally, we mention that there is a variant of the SOR method, called the
*Symmetric Successive Over-relaxation* method, known as SSOR, where
one runs a standard SOR sweep through the mesh points and then a
new sweep while visiting the points in reverse order.

## The Conjugate gradient method

There is no simple intuitive derivation of the Conjugate gradient
method, so we refer to the many excellent expositions in the
literature for the idea of the method and how the algorithm is
derived. In particular, we recommend the books
[[Templates_LA;@Axelsson_1996;@Saad_2003;@Greif_Ascher_2011]](#Templates_LA;@Axelsson_1996;@Saad_2003;@Greif_Ascher_2011).
A brief overview is
provided in the [Wikipedia article](https://en.wikipedia.org/wiki/Conjugate_gradient_method). Here, we
just state the pros and cons of the method from a user's perspective
and how we utilize it in code.

The original Conjugate gradient method is limited to linear systems $Au=b$,
where $A$ is a symmetric and positive definite matrix. There are, however,
extensions of the method to non-symmetric matrices.

A major advantage of all conjugate gradient methods is that the matrix
$A$ is only used in matrix-vector products, so we do not need form and
store $A$ if we can provide code for computing a matrix-vector product
$Au$.  Another important feature is that the algorithm is very easy to
vectorize and parallelize.  The primary downside of the method is that
it converges slowly unless one has an effective *preconditioner*
for the system. That is, instead of solving $Au=b$, we try to solve
$M^{-1}Au=M^{-1}b$ in the hope that the method works better for this
*preconditioned* system. The matrix $M$ is the *preconditioner* or
preconditioning matrix.  Now we need to perform matrix-vector products
$y = M^{-1}Au$, which is done in two steps: first the matrix-vector
product $v=Au$ is carried out and then the system $My=v$ must be
solved.  Therefore, $M$ must be cheap to compute and systems $My=v$
must be cheap to solve.

A perfect preconditioner is $M=A$, but in each iteration in the
Conjugate gradient method one then has so solve a system with $A$ as
coefficient matrix! A key idea is to let $M$ be some kind of *cheap
approximation* to $A$.  The simplest preconditioner is to set $M=D$,
where $D$ is the diagonal of $A$. This choice means running one Jacobi
iteration as preconditioner. [diffu:exer:splitting_prec](#diffu:exer:splitting_prec)
shows that the Jacobi and SOR methods can also be viewed as
preconditioners.

Constructing good preconditioners is a scientific field on its
own. Here we shall treat the topic just very briefly. For a user
having access to the `scipy.sparse.linalg` library, there are
Conjugate gradient methods and preconditioners readily available:

 * For positive definite, symmetric systems: `cg` (the Conjugate gradient method)

 * For symmetric systems: `minres` (Minimum residual method)

 * For non-symmetric systems:

  * `gmres` (GMRES: Generalized minimum residual method)

  * `bicg` (BiConjugate gradient method)

  * `bicgstab` (Stabilized BiConjugate gradient method)

  * `cgs` (Conjugate gradient squared method)

  * `qmr` (Quasi-minimal residual iteration)

 * Preconditioner: `spilu` (Sparse, incomplete LU factorization)

The ILU preconditioner is an attractive all-round type of
preconditioner that is suitable for most problems on serial
computers. A more efficient preconditioner is the multigrid method,
and algebraic multigrid is also an all-round choice as
preconditioner. The Python package [PyAMG](https://github.com/pyamg/pyamg) offers efficient implementations of
the algebraic multigrid method, to be used both as a preconditioner
and as a stand-alone iterative method.

The matrix arising from implicit time discretization methods applied to the
diffusion equation is symmetric and positive definite. Thus, we can use
the Conjugate gradient method (`cg`), typically in combination with an
ILU preconditioner. The code is very similar to the one we created
when solving the linear system by sparse Gaussian elimination, the
main difference is that we now allow for calling up the Conjugate
gradient function as an alternative solver.

```
def solver_sparse(
    I, a, f, Lx, Ly, Nx, Ny, dt, T, theta=0.5,
    U_0x=0, U_0y=0, U_Lx=0, U_Ly=0, user_action=None,
    method='direct', CG_prec='ILU', CG_tol=1E-5):
    """
    Full solver for the model problem using the theta-rule
    difference approximation in time. Sparse matrix with
    dedicated Gaussian elimination algorithm (method='direct')
    or ILU preconditioned Conjugate Gradients (method='CG' with
    tolerance CG_tol and preconditioner CG_prec ('ILU' or None)).
    """
    # Set up data structures as shown before

    # Precompute sparse matrix
    ...

    A = scipy.sparse.diags(
        diagonals=[main, lower, upper, lower2, upper2],
        offsets=[0, -lower_offset, lower_offset,
                 -lower2_offset, lower2_offset],
        shape=(N, N), format='csc')

    if method == 'CG':
        if CG_prec == 'ILU':
            # Find ILU preconditioner (constant in time)
            A_ilu = scipy.sparse.linalg.spilu(A)  # SuperLU defaults
            M = scipy.sparse.linalg.LinearOperator(
                shape=(N, N), matvec=A_ilu.solve)
        else:
            M = None
        CG_iter = []  # No of CG iterations at time level n

    # Time loop
    for n in mathcal{I}_t[0:-1]:
        # Compute b, vectorized version

        # Solve matrix system A*c = b
        if method == 'direct':
            c = scipy.sparse.linalg.spsolve(A, b)
        elif method == 'CG':
            x0 = u_n.T.reshape(N)  # Start vector is u_n
            CG_iter.append(0)

            def CG_callback(c_k):
                """Trick to count the no of iterations in CG."""
                CG_iter[-1] += 1

            c, info = scipy.sparse.linalg.cg(
                A, b, x0=x0, tol=CG_tol, maxiter=N, M=M,
                callback=CG_callback)

        # Fill u with vector c
        # Update u_n before next step
        u_n, u = u, u_n
```

The number of iterations in the Conjugate gradient method is of interest,
but is unfortunately not available from the `cg` function. Therefore,
we perform a trick: in each iteration a user function `CG_callback`
is called where we accumulate the number of iterations in a list
`CG_iter`.

## What is the recommended method for solving linear systems?

There is no clear answer to this question. If you have enough memory
and computing time available, direct methods such as `spsolve` are
to be preferred since they are easy to use and finds almost an exact
solution. However, in larger 2D and in 3D problems, direct methods
usually run too slowly or require too much memory, so one is forced
to use iterative methods. The fastest and most reliable methods are
in the Conjugate Gradient family, but these require suitable
preconditioners. ILU is an all-round preconditioner, but it is not
suited for parallel computing. The Jacobi and SOR iterative methods
are easy to implement, and popular for that reason, but run slowly. Jacobi
iteration is not an option in real problems, but SOR may be.


# [IndexedBase](https://stackoverflow.com/questions/37647370/expand-index-notation-equation-using-sympy)
> ### that is, symbols indexed by other symbols

# $$
\sigma_{ij} = \lambda U_{k,k}\delta_{ij} + \mu(U_{i,j} + U_{j,i})
$$
> ### comma($U_{k,k}$) is a convention for derivative
> ### $U_1, U_2, U_3$ are written as u,v,w
> ### $x_1, x_2, x_3$ are written as x,y,z

```python
import sympy as sm
U = sm.IndexedBase('U')
i,j,k,l,mu,x,y,z = sm.symbols('i j k lambda nu x y z')
u,v,w = sm.symbols('u v w',cls=sm.Function)

sij = l*sm.Sum(U[k,k],(k,0,2)) * sm.KroneckerDelta(i,j)
l*sm.Sum(U[k,k],(k,0,2)) * sm.KroneckerDelta(i,j) + mu*(U[i,j] + U[j,i])
## replace
du = sm.derive_by_array([u(x,y,z), v(x,y,z),w(x,y,z)],[x,y,z])
l*sm.tensorcontraction(du,(0,1))
```

$\displaystyle {U}_{i,j}$

### $$ \sigma_x = \lambda\Big(\frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} + \frac{\partial w}{\partial z} \Big)
+ 2\mu \frac{\partial u}{\partial x} $$
### $$ \sigma_y = \lambda\Big(\frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} + \frac{\partial w}{\partial z} \Big)
+ 2\mu \frac{\partial v}{\partial y} $$
### $$ \sigma_z = \lambda\Big(\frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} + \frac{\partial w}{\partial z} \Big)
+ 2\mu \frac{\partial w}{\partial z} $$
### $$
\tau_{xy}= \mu\Big( \frac{\partial u}{\partial y} + \frac{\partial v}{\partial x}\Big),
\\ \tau_{yz}= \mu\Big( \frac{\partial v}{\partial z} + \frac{\partial w}{\partial y}\Big),
\\ \tau_{zx}= \mu\Big( \frac{\partial w}{\partial x} + \frac{\partial u}{\partial z}\Big),
$$

```python
# U[0,0] -> Derivative(u,x)
def replace_U(uij):
    i1, i2 = uij,indices
    return sm.Derivative([u,v,w][i1],[x,y,z][i2])
for ii in range(3):
    for ji in range(3):
        if ii < ji:
            continue
        sm.pprint(sij.doit().xreplace({i:ii, j:ji}).replace(lambda v: v.base == U,replace_U))
```

```python

```


# Разработать программу, для решения численным методом гравитационной задачи N тел.

Эволюция системы N материальных точек описывается следующей системой уравнений:
\begin{equation}
    \begin{cases}
        \frac{dr_i}{dt} = v_i\\
        \frac{dv_i}{dt} = \displaystyle\sum_{i \ne j}^{N} Gm_j\frac{r_j - r_i}{|r_j - r_i|^3}
    \end{cases}
\end{equation}

```python
import numpy as np
import time
from scipy.constants import g
import matplotlib.pyplot as plt
from matplotlib.animation import ArtistAnimation
from IPython.display import HTML
```

#### Функция `getAcc` находит проекции ускорения на оси x и y.

```python
def getAcc(pos, mass, G, softening):
    x = pos[:, 0:1]
    y = pos[:, 1:2]

    dx = x.T - x
    dy = y.T - y

    inv_r3 = (dx ** 2 + dy ** 2 + softening ** 2) ** (-1.5)

    ax = G * (dx * inv_r3) @ mass
    ay = G * (dy * inv_r3) @ mass

    a = np.hstack((ax, ay))

    return a
```

### Задание начальных условий:

```python
t = 0
t_end = 10.0
dt = 0.01
softening = 0.1
field_size = 500
```

```python
n = int(input("Введите количество частиц: "))
des_1 = int(input("Введите 1, если масса различна, иначе 2: "))
des_2 = int(input("Введите 1, если первоначальная скорость = 0, иначе 2: "))

```

    Введите количество частиц: 50
    Введите 1, если масса различна, иначе 2: 1
    Введите 1, если первоначальная скорость = 0, иначе 2: 2

```python
if des_1 == 1:
    mass = np.random.random((n, 1)).astype(float) * 10 ** 4
else:
    mass = np.ones((n, 1)) * 2 * 10**4

if des_2 == 1:
    vel = np.zeros((n, 2))
else:
    vel = np.random.randn(n, 2)

pos = np.random.randn(n, 2) * 100
```

```python
acc = getAcc(pos, mass, g, softening)
```

```python
color = []
pos_x = {}
pos_y = {}
for i in range(n):
    pos_x[i] = []
    pos_y[i] = []
    color.append((
        np.random.random(),
        np.random.random(),
        np.random.random()
    ))
```

### Построение модели

#### Переменная `step_count` содержит количетсво шагов моделирования.

```python
step_count = int(np.ceil(t_end / dt))
```

```python
fig, ax = plt.subplots()

ax.set_xlim((-field_size, field_size))
ax.set_ylim((-field_size, field_size))

plt.close()
frames = []
```

### Осуществление вычислений

Код просчитывает ускорение и скорость для всех заданных точек.

А также сохраняет кадры для анимации.

```python
for i in range(step_count):
    vel += acc * dt / 2.0
    pos += vel * dt
    line = []
    for j in range(len(pos)):
        pos_x[j].append(pos[j][0])
        pos_y[j].append(pos[j][1])
        temp, = ax.plot(pos_x[j], pos_y[j], color=color[j], linewidth=1)
        line.append(temp,)
    frames.append(line)
    acc = getAcc(pos, mass, g, softening)
    vel += acc * dt / 2.0
    t += dt
```

Создать анимацию:

```python
anim = ArtistAnimation(
    fig,
    frames,
    interval=60,
    blit=True,
    repeat=True
)
```

Отобразить анимацию.

```python
HTML(anim.to_html5_video())
```

individual tasks/N Body simulation.ipynb