Skip to main content

qualia_core_db/solvers/interpolation/
least_squares.rs

1//! Polynomial least-squares fitting via the normal equations `(VᵀV) c = Vᵀy`, solved by
2//! Gaussian elimination with partial pivoting (the small `(d+1)×(d+1)` system).
3
4use super::InterpolationError;
5
6/// Fit a degree-`degree` polynomial to `(xs, ys)` in the least-squares sense. Returns
7/// coefficients in **ascending** order `[c₀, c₁, …, c_degree]` (so the polynomial is
8/// `Σ cₖ xᵏ`). Fails closed if there are too few points or the system is singular.
9pub fn poly_fit(xs: &[f64], ys: &[f64], degree: usize) -> Result<Vec<f64>, InterpolationError> {
10    if xs.is_empty() || xs.len() != ys.len() {
11        return Err(InterpolationError::InsufficientData);
12    }
13    if degree + 1 > xs.len() {
14        return Err(InterpolationError::InvalidDegree);
15    }
16    let m = degree + 1;
17    // Normal equations: A[j][k] = Σ x^(j+k), b[j] = Σ y·x^j.
18    let mut a = vec![0.0; m * m];
19    let mut b = vec![0.0; m];
20    // Precompute power sums up to 2·degree.
21    let mut powsum = vec![0.0; 2 * degree + 1];
22    for &x in xs {
23        let mut p = 1.0;
24        for s in powsum.iter_mut() {
25            *s += p;
26            p *= x;
27        }
28    }
29    for j in 0..m {
30        for k in 0..m {
31            a[j * m + k] = powsum[j + k];
32        }
33        let mut s = 0.0;
34        for (&x, &y) in xs.iter().zip(ys) {
35            s += y * x.powi(j as i32);
36        }
37        b[j] = s;
38    }
39    gauss_solve(m, &mut a, &mut b).ok_or(InterpolationError::Singular)
40}
41
42/// Evaluate a polynomial given ascending coefficients at `x` (Horner).
43pub fn poly_eval(coeffs: &[f64], x: f64) -> f64 {
44    coeffs.iter().rev().fold(0.0, |acc, &c| acc * x + c)
45}
46
47/// Solve `A x = b` (row-major `n×n`) by Gaussian elimination with partial pivoting.
48/// Consumes `a`/`b`. `None` if singular.
49fn gauss_solve(n: usize, a: &mut [f64], b: &mut [f64]) -> Option<Vec<f64>> {
50    for col in 0..n {
51        // Partial pivot.
52        let mut piv = col;
53        let mut best = a[col * n + col].abs();
54        for r in (col + 1)..n {
55            let v = a[r * n + col].abs();
56            if v > best {
57                best = v;
58                piv = r;
59            }
60        }
61        if best < 1e-14 {
62            return None; // singular
63        }
64        if piv != col {
65            for c in 0..n {
66                a.swap(col * n + c, piv * n + c);
67            }
68            b.swap(col, piv);
69        }
70        // Eliminate below.
71        for r in (col + 1)..n {
72            let f = a[r * n + col] / a[col * n + col];
73            for c in col..n {
74                a[r * n + c] -= f * a[col * n + c];
75            }
76            b[r] -= f * b[col];
77        }
78    }
79    // Back-substitution.
80    let mut x = vec![0.0; n];
81    for i in (0..n).rev() {
82        let mut s = b[i];
83        for j in (i + 1)..n {
84            s -= a[i * n + j] * x[j];
85        }
86        x[i] = s / a[i * n + i];
87    }
88    Some(x)
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn recovers_a_line_from_collinear_points() {
97        // y = 3x + 1
98        let xs = [0.0, 1.0, 2.0, 3.0, 4.0];
99        let ys = xs.map(|x| 3.0 * x + 1.0);
100        let c = poly_fit(&xs, &ys, 1).unwrap();
101        assert!((c[0] - 1.0).abs() < 1e-9);
102        assert!((c[1] - 3.0).abs() < 1e-9);
103    }
104
105    #[test]
106    fn recovers_a_parabola_exactly() {
107        // y = 2x² − x + 5, fit degree 2 → exact.
108        let f = |x: f64| 2.0 * x * x - x + 5.0;
109        let xs = [-2.0, -1.0, 0.0, 1.0, 2.0, 3.0];
110        let ys = xs.map(f);
111        let c = poly_fit(&xs, &ys, 2).unwrap();
112        for &q in &[0.5, 1.3, -1.7] {
113            assert!((poly_eval(&c, q) - f(q)).abs() < 1e-7);
114        }
115    }
116
117    #[test]
118    fn least_squares_minimises_on_noisy_data() {
119        // Points near y = x; degree-1 fit slope ≈ 1, intercept ≈ 0.
120        let xs = [0.0, 1.0, 2.0, 3.0, 4.0];
121        let ys = [0.1, 0.9, 2.1, 2.9, 4.05];
122        let c = poly_fit(&xs, &ys, 1).unwrap();
123        assert!((c[1] - 1.0).abs() < 0.1);
124        assert!(c[0].abs() < 0.2);
125    }
126
127    #[test]
128    fn fails_closed() {
129        assert_eq!(
130            poly_fit(&[1.0, 2.0], &[1.0, 2.0], 5).unwrap_err(),
131            InterpolationError::InvalidDegree
132        );
133        assert_eq!(
134            poly_fit(&[], &[], 0).unwrap_err(),
135            InterpolationError::InsufficientData
136        );
137    }
138}