Skip to main content

solver/linalg/
adi.rs

1/// Thomas algorithm (tridiagonal matrix algorithm) for solving a 1-D tridiagonal system.
2///
3/// Solves: `lower[i]*x[i-1] + diag[i]*x[i] + upper[i]*x[i+1] = rhs[i]`
4///
5/// Boundary conventions:
6/// - `lower[0]` is unused (no element below the first row).
7/// - `upper[n-1]` is unused (no element above the last row).
8///
9/// # Arguments
10/// * `n`     - System size.
11/// * `lower` - Sub-diagonal coefficients (length `n`; `lower[0]` ignored).
12/// * `diag`  - Main diagonal coefficients (length `n`).
13/// * `upper` - Super-diagonal coefficients (length `n`; `upper[n-1]` ignored).
14/// * `rhs`   - Right-hand side vector (length `n`).
15/// * `x`     - Output solution vector (length `n`).
16///
17/// # Panics
18/// Panics in debug mode if `n == 0`.
19///
20/// # Examples
21///
22/// ```
23/// use solver::linalg::adi::thomas;
24///
25/// // Solve the system:
26/// //  [ 2 -1  0 ] [x0]   [1]
27/// //  [-1  2 -1 ] [x1] = [0]
28/// //  [ 0 -1  2 ] [x2]   [1]
29/// // Solution: x = [1, 1, 1]
30/// let lo  = [0.0, -1.0, -1.0];
31/// let di  = [2.0,  2.0,  2.0];
32/// let up  = [-1.0, -1.0, 0.0];
33/// let rhs = [1.0,  0.0,  1.0];
34/// let mut x = [0.0f64; 3];
35/// thomas(3, &lo, &di, &up, &rhs, &mut x);
36/// assert!((x[0] - 1.0).abs() < 1e-12);
37/// assert!((x[1] - 1.0).abs() < 1e-12);
38/// assert!((x[2] - 1.0).abs() < 1e-12);
39/// ```
40pub fn thomas(n: usize, lower: &[f64], diag: &[f64], upper: &[f64], rhs: &[f64], x: &mut [f64]) {
41    debug_assert!(n > 0, "thomas: system size must be > 0");
42    let mut c = vec![0.0f64; n];
43    let mut d = vec![0.0f64; n];
44
45    // Forward sweep
46    let w0 = diag[0];
47    c[0] = if n > 1 { upper[0] / w0 } else { 0.0 };
48    d[0] = rhs[0] / w0;
49    for i in 1..n {
50        let w = diag[i] - lower[i] * c[i - 1];
51        c[i] = if i < n - 1 { upper[i] / w } else { 0.0 };
52        d[i] = (rhs[i] - lower[i] * d[i - 1]) / w;
53    }
54
55    // Back substitution
56    x[n - 1] = d[n - 1];
57    for i in (0..n - 1).rev() {
58        x[i] = d[i] - c[i] * x[i + 1];
59    }
60}