Tridiagonal Solver Example — Thomas Algorithm

Demonstrates the Thomas algorithm via linearsystems.tridiag_solve() by solving the 1D Poisson equation \(-u'' = \pi^2 \sin(\pi x)\) on \([0,1]\) with Dirichlet boundary conditions, whose exact solution is \(u(x) = \sin(\pi x)\). The result is compared against a full-matrix solve using linearsystems.lslrg().

Example Code

"""IMSL Tridiagonal Solver example: 1D Poisson equation with Thomas algorithm.

Solves -u'' = f(x) on [0,1] with Dirichlet BCs u(0)=u(1)=0 and
f(x) = π²sin(πx), whose exact solution is u(x)=sin(πx).

Discretizing on an n-point interior grid gives a tridiagonal system. The
Thomas algorithm (tridiag_solve) is compared against a full-matrix solve
(lslrg) for accuracy and speed.

Outputs:
- Max error and solve times printed to stdout
- SVG plot of numerical vs analytical solution saved to
  test_output/example_imsl_tridiagonal.svg
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Dict

import matplotlib.pyplot as plt
import numpy as np

from linearsystems import tridiag_solve, lslrg


def run_demo_imsl_tridiagonal() -> Dict[str, object]:
    """Solve the 1D Poisson equation using tridiag_solve (Thomas algorithm).

    Discretises -u'' = π²sin(πx) on n=100 interior points, forms the
    tridiagonal system, solves with tridiag_solve(), and also with lslrg()
    for comparison. Reports max absolute error against sin(πx).

    Args:
        None

    Returns:
        Dict[str, object]: Result dict with keys ``max_error_tridiag``,
            ``max_error_full``, ``time_tridiag``, ``time_full``,
            and ``plot_path`` (str).
    """
    n = 100  # interior grid points
    h = 1.0 / (n + 1)
    x = np.linspace(h, 1.0 - h, n)       # interior nodes
    x_full = np.linspace(0, 1, n + 2)     # including BCs for plotting

    # Right-hand side: f(x) = π²sin(πx)
    f = np.pi**2 * np.sin(np.pi * x)

    # Exact solution
    u_exact = np.sin(np.pi * x)

    # Build tridiagonal system: (2/h²)u_i - (1/h²)u_{i-1} - (1/h²)u_{i+1} = f_i
    diag_main = np.full(n, 2.0 / h**2)
    diag_off = np.full(n - 1, -1.0 / h**2)

    # Thomas algorithm solve
    t0 = time.perf_counter()
    u_tridiag = tridiag_solve(diag_off, diag_main, diag_off, f)
    t_tridiag = time.perf_counter() - t0

    # Full-matrix solve (lslrg) for comparison
    a_full = np.diag(diag_main) + np.diag(diag_off, 1) + np.diag(diag_off, -1)
    t0 = time.perf_counter()
    result_full = lslrg(a_full, f)
    t_full = time.perf_counter() - t0
    u_full = result_full.x

    max_err_tridiag = float(np.max(np.abs(u_tridiag - u_exact)))
    max_err_full = float(np.max(np.abs(u_full - u_exact)))

    print("\nIMSL Tridiagonal Solver: 1D Poisson Equation (-u'' = π²sin(πx))")
    print("-" * 65)
    print(f"Grid points (interior): n = {n},  h = {h:.6f}")
    print(f"{'Method':<30} {'Max |error|':>14} {'Solve time (s)':>16}")
    print("-" * 65)
    print(f"{'tridiag_solve (Thomas)':<30} {max_err_tridiag:>14.4e} {t_tridiag:>16.6f}")
    print(f"{'lslrg (full matrix)':<30} {max_err_full:>14.4e} {t_full:>16.6f}")
    print("-" * 65)
    print(f"\nExpected O(h²) error ≈ {h**2:.4e}  (h²π²/12 estimate)")

    output_dir = Path("test_output")
    output_dir.mkdir(parents=True, exist_ok=True)
    plot_path = output_dir / "example_imsl_tridiagonal.svg"

    # Embed BCs (u=0) for plotting
    u_plot_tridiag = np.concatenate([[0.0], u_tridiag, [0.0]])
    u_plot_exact = np.concatenate([[0.0], u_exact, [0.0]])
    error_plot = np.concatenate([[0.0], u_tridiag - u_exact, [0.0]])

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Left: solution comparison
    ax = axes[0]
    ax.plot(x_full, u_plot_exact, "k-", linewidth=2.5, label="Exact: sin(πx)")
    ax.plot(x_full, u_plot_tridiag, "b--", linewidth=1.5, label="tridiag_solve (Thomas)")
    ax.set_xlabel("x")
    ax.set_ylabel("u(x)")
    ax.set_title("1D Poisson: Numerical vs Analytical")
    ax.legend()
    ax.grid(True, alpha=0.3)

    # Right: pointwise error
    ax2 = axes[1]
    ax2.plot(x_full, error_plot, "r-", linewidth=1.5)
    ax2.axhline(0, color="k", linewidth=0.8, linestyle="--")
    ax2.set_xlabel("x")
    ax2.set_ylabel("u_numerical − u_exact")
    ax2.set_title(f"Pointwise Error  (max={max_err_tridiag:.2e})")
    ax2.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(plot_path, format="svg")
    plt.close(fig)
    print(f"\nPlot saved to: {plot_path}")

    return {
        "max_error_tridiag": max_err_tridiag,
        "max_error_full": max_err_full,
        "time_tridiag": t_tridiag,
        "time_full": t_full,
        "plot_path": str(plot_path),
    }


if __name__ == "__main__":
    run_demo_imsl_tridiagonal()

Plot Output

Left: numerical solution (Thomas algorithm) overlaid on the exact \(\sin(\pi x)\) curve. Right: pointwise error showing the expected \(O(h^2)\) spatial convergence.

Console Output

IMSL Tridiagonal Solver: 1D Poisson Equation (-u'' = π²sin(πx))
-----------------------------------------------------------------
Grid points (interior): n = 100,  h = 0.009901
Method                            Max |error|   Solve time (s)
-----------------------------------------------------------------
tridiag_solve (Thomas)             8.0620e-05         0.000174
lslrg (full matrix)                8.0620e-05         0.026109
-----------------------------------------------------------------

Expected O(h²) error ≈ 9.8030e-05  (h²π²/12 estimate)

Plot saved to: test_output\example_imsl_tridiagonal.svg