solve-block-tridiagonal.py

import sys
import numpy as np
import numpy.linalg as la
from scipy.io import loadmat
 
# Load the data generated by solve-block-tridiagonal.cpp
data = loadmat("block_tridiagonal_system.mat")
# Data is stored in Fortran order, so transpose
M = np.moveaxis(data["M"], -1, 0)  # (n, n, N) -> (N, n, n)
K = np.moveaxis(data["K"], -1, 0)  # (n, n, N) -> (N, n, n)
b = np.moveaxis(data["b"], -1, 0)  # (n, 1, N) -> (N, n, 1)
x = np.moveaxis(data["x"], -1, 0)  # (n, 1, N) -> (N, n, 1)
N, n, _ = M.shape
 
# Construct the full block tridiagonal system as a dense matrix for verification
A = np.zeros((N * n, N * n))
for i in range(N):
    A[i * n : (i + 1) * n, i * n : (i + 1) * n] = np.tril(M[i]) + np.tril(M[i], -1).T
    A[((i + 1) % N) * n : ((i + 1) % N + 1) * n, i * n : (i + 1) * n] = K[i]
    A[i * n : (i + 1) * n, ((i + 1) % N) * n : ((i + 1) % N + 1) * n] = K[i].T
 
# Report the condition number and the residual of the solution
print(f"Number of blocks:    {N}")
print(f"Block size:          {n}")
print(f"Number of variables: {N * n}")
print(f"Number of nonzeros:  {N * n * (n + 1) // 2 + N * n * n}")
σA = la.eigvalsh(A)
nrmA = max(σA)
condA = nrmA / min(σA)
res = la.norm(A @ x.ravel() - b.ravel())
nrmx, nrmb = la.norm(x.ravel()), la.norm(b.ravel())
err = res / (nrmA * nrmx + nrmb)
print(f"Condition number κ(A) =         {condA:.17e}")
print(f"Relative residual ‖Ax-b‖/‖b‖ =  {res / nrmb:.17e}")
print(f"                    ‖Ax-b‖")
print(f"Backward error η = ────────── = {err:.17e}")
print(f"                   ‖A‖‖x‖+‖b‖")
 
sys.exit(0 if err < 10 * np.finfo(M.dtype).eps * N * n else 1)