rustyqlib/core/fd_solvers/
tridiagonal.rs

1//! The Thomas algorithm: direct O(n) solve of a tridiagonal system.
2
3/// Solves a tridiagonal system `A x = d` where `a` is the sub-diagonal
4/// (`a[i-1]` multiplies `x[i-1]` in row `i`), `b` the diagonal and `c` the
5/// super-diagonal (`c[i]` multiplies `x[i+1]` in row `i`).
6/// https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
7pub fn thomas_algorithm(a: &[f64], b: &[f64], c: &[f64], d: &[f64]) -> Vec<f64> {
8    let n = d.len();
9    assert!(b.len() == n && a.len() == n - 1 && c.len() == n - 1);
10    if n == 1 {
11        return vec![d[0] / b[0]];
12    }
13    let mut c_ = c.to_vec();
14    let mut d_ = d.to_vec();
15    let mut x: Vec<f64> = vec![0.0; n];
16
17    c_[0] = c_[0] / b[0];
18    d_[0] = d_[0] / b[0];
19    for i in 1..n - 1 {
20        let id = 1.0 / (b[i] - a[i - 1] * c_[i - 1]);
21        c_[i] = c_[i] * id;
22        d_[i] = (d_[i] - a[i - 1] * d_[i - 1]) * id;
23    }
24    d_[n - 1] = (d_[n - 1] - a[n - 2] * d_[n - 2]) / (b[n - 1] - a[n - 2] * c_[n - 2]);
25
26    x[n - 1] = d_[n - 1];
27    for i in (0..n - 1).rev() {
28        x[i] = d_[i] - c_[i] * x[i + 1];
29    }
30    x
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36
37    #[test]
38    fn thomas_solves_small_system() {
39        // [2 1 0; 1 2 1; 0 1 2] x = [4; 8; 8] -> x = [1; 2; 3]
40        let x = thomas_algorithm(&[1.0, 1.0], &[2.0, 2.0, 2.0], &[1.0, 1.0], &[4.0, 8.0, 8.0]);
41        for (got, want) in x.iter().zip(&[1.0, 2.0, 3.0]) {
42            assert!((got - want).abs() < 1e-12, "{x:?}");
43        }
44    }
45}
rustyqlib/core/fd_solvers/tridiagonal.rs

rustyqlib/core/fd_solvers/
tridiagonal.rs
