Saturday 17 August 2013

Three-diagonal system of linear equations

In this post I'll show a fast and accurate solution for the three-diagonal system of linear equations based on Gaussian elimination. In the matrix notation the system has the following form:

\begin{equation}\left(\begin{matrix} b_0 & c_0 & 0   & \dots & 0   & 0  & 0 \\
a_1 & b_1 & c_1 & \dots & 0   & 0  & 0\\
\vdots & \vdots & \vdots & \ddots & \vdots   & \vdots  & \vdots  \\
0 & 0 & 0 & \dots & a_{n-1} & b_{n-1} & c_{n-1}  \\
0 & 0 & 0 & \dots & 0 & a_{n} & b_n \end{matrix}\right)
\left(\begin{matrix} x_0 \\ x_1 \\ \vdots \\ x_{n-1}  \\x_{n}\end{matrix}\right) = \left(\begin{matrix} d_0 \\ d_1 \\ \vdots \\ d_{n-1}  \\d_{n}\end{matrix}\right)
\label{eq:matrix00}\nonumber\end{equation}

or in the form of the individual equations ($n+1$ equation with $n+1$ unknown $x_i$):

\begin{eqnarray}
b_0 x_0 + c_0 x_1 &=& d_0,\nonumber\\
a_1 x_0 + b_1 x_1 + c_1 x_2 &=& d_1,\nonumber\\
\dots &=& \dots, \nonumber\\
a_i x_{i-1}+ b_i x_i + c_i x_{i+1} &=& d_i,\nonumber\\
\dots &=& \dots ,\nonumber\\
a_{n-1} x_{n-2} +b_{n-1} x_{n-1} + c_{n-1} x_{n} &=& d_{n-1},\nonumber\\
a_{n} x_{n-1} +b_{n} x_{n} &=& d_{n}.\nonumber
\end{eqnarray}


The system has three diagonals. The main diagonal is populated with $n+1$ coefficients $b_i$. The side diagonals are populated by $n$ coefficients $a$ (below, indices from 1 to n) and $n$ coefficients $c$ (above, indices from 0 to n-1).

The solution. The unknown $x_0$ can be expressed using the first equation:
$$
x_0 = (d_0 - c_0 x_1) \frac{1}{b_0}.
$$

With that expression we can substitute $x_0$ in the second equation:

$$
(d_0 - c_0 x_1) \frac{a_1}{b_0} + b_1 x_1 + c_1 x_2 = d_1,
$$
$$
\left(b_1 - \frac{a_1 c_0}{b_0}\right) x_1 + c_1 x_2 = \left(d_1 - \frac{a_1 d_0}{b_0}\right),
$$
$$
b_1' x_1 + c_1 x_2 = d_1',
$$
where
$$
b_1' = b_1 - \frac{a_1 c_0}{b_0},
$$
and
$$
d_1' = d_1 - \frac{a_1 d_0}{b_0}.
$$

Now the second equation has exactly the same form as the first one, so we can use it to express $x_1$ that we are going to substitute into the next equation:
$$
b_2' x_2 + c_2 x_3 = d_2',
$$
where
$$
b_2' = b_2 - \frac{a_2 c_1}{b_1'},
$$
and
$$
d_2' = d_2 - \frac{a_2 d_1}{b_1'}.
$$

And so on. The i-th equation becomes
$$
b_i' x_i + c_i x_{i+1} = d_i',
$$
where
$$
b_i' = b_i - \frac{a_i c_{i-1}}{b_{i-1}'},
$$
and
$$
d_i' = d_i - \frac{a_i d_{i-1}}{b_{i-1}'}.
$$
At the end of the sequence the n-1-st equation becomes:
$$
b_{n-1}' x_{n-1} + c_{n-1} x_{n} = d_{n-1}',
$$
with
$$
b_{n-1}' = b_{n-1} - \frac{a_{n-1} c_{n-2}}{b_{n-2}'},
$$
and
$$
d_{n-1}' = d_{n-1} - \frac{a_{n-1} d_{n-2}}{b_{n-2}'}.
$$
Finally, in the last equation we can substitute $x_{n-1}$ as
$$
x_{n-1} = (d_{n-1}' - c_{n-1} x_{n})\frac{1}{b_{n-1}' },
$$
so the last equation becomes a linear equation only of $x_n$:
$$
(d_{n-1}' - c_{n-1} x_{n})\frac{a_n}{b_{n-1}' } +b_{n} x_{n} = d_{n},
$$
$$
\left(b_n - \frac{a_n c_{n-1}}{b_{n-1}' }\right) x_n = d_{n} - \frac{a_n d_{n-1}'}{b_{n-1}' },
$$
and simply:
$$
b_n' x_n = d_n'
$$
with
$$
b_{n-1}' = b_n - \frac{a_n c_{n-1}}{b_{n-1}' },
$$
and
$$
d_{n-1}' = d_{n} - \frac{a_n d_{n-1}'}{b_{n-1}' }.
$$
The solution for $x_n$ is clearly:
$$
x_n = \frac{d_n'}{b_n'}.
$$
Now we can do back substitutions to get other unknowns $x_i$ for $i = 0, \dots, n-1$:
\begin{eqnarray}
x_{n-1} &=& (d_{n-1}' - c_{n-1} x_{n})\frac{1}{b_{n-1}' },\nonumber\\
\dots &=& \dots,\nonumber\\
x_{i} &=& (d_{i}' - c_{i} x_{i+1})\frac{1}{b_{i}' },\nonumber\\
\dots &=& \dots,\nonumber\\
x_{0} &=& (d_{0}' - c_{0} x_{1})\frac{1}{b_{0}' }.\nonumber
\end{eqnarray}

Implementation of this algorithm in any programming language is almost trivial. My IDL code that solves the problem is in the threediagonal_system.pro function. As the input the code takes just 4 vectors: a, b, c and d with the meaning as in this text above. Note that all 4 arrays have to be of the same size. The first element of the a array and the last element of the c array must be equal to 0.

Example:
; We form a system of five equations
IDL> aa = [0., 2.,  3., 4., 1.]
IDL> bb = [3., 4., 11., 7., 2.]
IDL> cc = [1., 1.,  1., 3., 0.]

; Set the exact solution
IDL> x0 = FINDGEN(5)

; Compute the right-hand side to complete the system
IDL> dd = FLTARR(5)
IDL> dd[0] = bb[0]*x0[0] + cc[0]*x0[1] + aa[0]*x0[4]
IDL> FOR i = 1, 3 DO dd[i] = aa[i]*x0[i-1] + bb[i]*x0[i] + cc[i]*x0[i+1]
IDL> dd[4] = aa[4]*x0[3] + bb[4]*x0[4] + cc[4]*x0[0]

; Solve the system
IDL> x = threediagonal_system(aa, bb, cc, dd)

; Compare the solution of the system and the exact solution
IDL> PRINT, x
 -3.97364e-08      1.00000      2.00000      3.00000      4.00000
IDL> PRINT, x0
      0.00000      1.00000      2.00000      3.00000      4.00000

No comments:

Post a Comment