import numpy as np
def solve_tridiagonal(U, L, D, F):
"""
Solve a tridiagonal system using the double-sweep method.
Parameters:
-----------
U : Upper diagonal elements (U[-1] is not used)
L : Lower diagonal elements (L[0] is not used)
D : Main diagonal elements
F : Right hand side vector
Returns:
--------
x : ndarray
Solution vector
Raises:
-------
ValueError
If the coefficients don't satisfy the conditions from equation (y23):
L_i, U_i > 0, D_i < 0, and -D_i ≥ L_i + U_i
"""
n = len(F)
# Check conditions from equation (y23)
# Skip U[-1] and L[0] as they are not used
if not np.all(U[:-1] > 0):
raise ValueError("Condition U_i > 0 not satisfied")
if not np.all(L[1:] > 0):
raise ValueError("Condition L_i > 0 not satisfied")
if not np.all(D < 0):
raise ValueError("Condition D_i < 0 not satisfied")
# Check -D_i ≥ U_i + L_i for i=1,...,n-1
# For i=1: only need U[0] since L[1] is the first lower diagonal element
if -D[0] < U[0]:
raise ValueError("Condition -D_1 ≥ U_1 not satisfied")
# For i=n: only need L[n] since U[n-1] is the last upper diagonal element
if -D[-1] < L[-1]:
raise ValueError("Condition -D_n ≥ L_n not satisfied")
# For i=2,...,n-1: need both L[i] and U[i-1]
if not np.all(-D[1:-1] >= L[2:] + U[:-2]):
raise ValueError("Condition -D_i ≥ L_i + U_i not satisfied")
alpha = np.zeros(n + 1)
beta = np.zeros(n + 1)
x = np.zeros(n + 1)
# First sweep: forward
# Initial conditions: alpha[1] = beta[1] = 0
for i in range(n):
denominator = D[i] + alpha[i] * L[i]
alpha[i+1] = -U[i] / denominator
beta[i+1] = (F[i] - L[i] * beta[i]) / denominator
# Second sweep: backward
# Initial condition: x[n+1] = 0
for i in range(n-1, -1, -1):
x[i] = alpha[i+1] * x[i+1] + beta[i+1]
return x[0:n]9 A direct method for solving tridiagonal linear systems
Consider the following system of linear equations \[ A\mathbf{x}= \mathbf{F}, \tag{9.1}\]
where \(\mathbf{F}\in\mathbb{R}^n\) is a given vector, \(\mathbf{x}\in\mathbb{R}^n\) is the vector of unknowns and \(A\) is a given \(n\times n\) tridiagonal matrix, i.e. \[ A= \begin{pmatrix} D_1 &U_1 &0 &\cdots &\cdots &\cdots &0 \\ L_2 &D_2 &U_2 &\ddots & & &\vdots \\ 0 &L_3 &D_3 &U_3 &\ddots & &\vdots \\ \vdots &\ddots &\ddots &\ddots &\ddots &\ddots &\vdots \\ \vdots & &\ddots &\ddots &\ddots &\ddots &0 \\ \vdots & & &\ddots &\ddots &D_{n-1} &U_{n-1} \\ 0 &\cdots &\cdots &\cdots &0 &L_n &D_n \end{pmatrix}. \tag{9.2}\]
A system like this has appeared when we discussed spline interpolation (see Eq. 6.70). Similar linear systems also arise in finite-difference methods for solving differential equations. Of course, the solution of this system can be approximated using iterative techniques such as Jacobi or Gauss-Seidel methods. However, there are much more efficient direct methods for solving such systems. All direct methods are equivalent to Gaussian elimination applied to the above trigiagonal system. One of these, called the double-sweep method, is described below.
The above system of linear equations can be written as \[ \begin{split} D_{1}x_{i}+U_{1}x_{2}&=F_{1}, \\ L_{i}x_{i-1}+D_{i}x_{i}+U_{i}x_{i+1}&=F_{i} \quad \hbox{for}\quad i=2, \dots, n-1,\\ L_{n}x_{n-1}+D_{n}x_{n}&=F_{n}. \end{split} \tag{9.3}\]
It is convenient to introduce \[ x_0=0 \quad \hbox{and} \quad x_{n+1}=0. \tag{9.4}\]
Then Eq. 9.3 can be rewritten as \[ L_{i}x_{i-1}+D_{i}x_{i}+U_{i}x_{i+1}=F_{i} \quad \hbox{for}\quad i=1, \dots, n. \tag{9.5}\]
To solve Eq. 9.5, we will seek \(\alpha_{i}\) and \(\beta_{i}\) such that \[ x_{i-1}=\alpha_{i}x_{i}+\beta_{i} \quad \hbox{for} \quad i=1, 2, \dots, n+1. \tag{9.6}\] Substitution of Eq. 9.6 into Eq. 9.5 yields \[ (\alpha_{i}L_{i}+D_{i})x_{i}+U_{i}x_{i+1}+\beta_{i}L_{i}-F_{i}=0 \quad \hbox{for}\quad i=1, \dots, n. \tag{9.7}\] From Eq. 9.6, we also have \[ x_{i}=\alpha_{i+1}x_{i+1}+\beta_{i+1} \quad \hbox{for} \quad i=0, 1, \dots, n. \tag{9.8}\] Substituting this into Eq. 9.7, we find that \[ [(\alpha_{i}L_{i}+D_{i})\alpha_{i+1}+U_{i}]x_{i+1}+[ (\alpha_{i}L_{i}+D_{i})\beta_{i+1}+\beta_{i}L_{i}-F_{i}]=0 \quad \hbox{for}\quad i=1, \dots, n. \tag{9.9}\] The last equation is satisfied if the two expressions in the square brackets are both zero. This leads to the following recursive formulae: \[ \alpha_{i+1}=-\frac{U_{i}}{D_{i}+\alpha_{i}L_{i}}, \quad \beta_{i+1}=-\frac{\beta_{i}L_{i}-F_{i}}{D_{i}+\alpha_{i}L_{i}}, \quad \hbox{for}\quad i=1, \dots, n. \tag{9.10}\] Now if \(\alpha_{1}\) and \(\beta_{1}\) are known, then \(\alpha_{i}\) and \(\beta_{i}\) for \(i=2, 3, \dots, n+1\) can be computed from Eq. 9.10. \(\alpha_{1}\) and \(\beta_{1}\) can be determined from Eq. 9.6 and the fact that \(x_{0}=0\). Indeed, \[ x_{0}= \alpha_{1}x_{1}+\beta_{1} \quad \hbox{and} \quad x_{0}=0 \quad \Rightarrow \quad \alpha_{1}x_{1}+\beta_{1}=0. \tag{9.11}\] To satisfy the last equation, we choose \(\alpha_{1}=0\) and \(\beta_{1}=0\). Once we know all \(\alpha_{i}\) and \(\beta_{i}\), we compute \(x_{n}, x_{n-1}, \dots, x_{1}\) using formula Eq. 9.6.
Formulae Eq. 9.6 and Eq. 9.10 will work provided that the coefficients \(L_{i}\), \(U_{i}\) and \(D_{i}\) are such that \(D_{i}+\alpha_{i}L_{i}\neq 0\) for \(i=1,\dots,n\). For tridiagonal systems that arise in finite-difference methods for differential equations, the coefficients \(L_{i}\), \(U_{i}\) and \(D_{i}\) usually satisfy the inequalities \[ L_{i}, U_{i} > 0, \quad D_{i} < 0, \quad -D_{i} \geq L_{i} + U_{i}. \tag{9.12}\] It can be shown that these restrictions on \(L_{i}\), \(U_{i}\) and \(D_{i}\) are sufficient for the double-sweep method to work.
The following function implements this method in Python:
Let us look at a simple example of a tri-diagonal system:
L = np.array([0, 1, 1, 1, 1]) # lower diagonal (L₂ to Lₙ)
U = np.array([1, 1, 1, 1, 0]) # upper diagonal (U₁ to Uₙ₋₁)
D = np.array([-2, -2, -2, -2, -2]) # main diagonal (D₁ to Dₙ)
F = np.array([1, 0, 0, 0, 1]) # right-hand side
# Solve using our function
x = solve_tridiagonal(U, L, D, F)
print("Solution:", x)Solution: [-1. -1. -1. -1. -1.]Let’s check that \(x\) does indeed satisfy Eq. 9.1. First we need to construct the full matrix \(A\).
# First, construct the full matrix
n = L.size
A = np.zeros((n, n))
for i in range(n):
if i < n-1:
A[i, i+1] = U[i] # upper diagonal
A[i, i] = D[i] # main diagonal
if i > 0:
A[i, i-1] = L[i] # lower diagonal
Aarray([[-2., 1., 0., 0., 0.],
[ 1., -2., 1., 0., 0.],
[ 0., 1., -2., 1., 0.],
[ 0., 0., 1., -2., 1.],
[ 0., 0., 0., 1., -2.]])Now we can use Python’s matrix multiplication opeator @ to calculate \(A\mathbf{x}\):
A @ xarray([ 1.00000000e+00, -2.22044605e-16, 0.00000000e+00, 0.00000000e+00,
1.00000000e+00])This indeed agrees with \(\mathbf{F}\) up to rounding errors.