qualia_core_db/solvers/interpolation/
least_squares.rs1use super::InterpolationError;
5
6pub 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 let mut a = vec![0.0; m * m];
19 let mut b = vec![0.0; m];
20 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
42pub fn poly_eval(coeffs: &[f64], x: f64) -> f64 {
44 coeffs.iter().rev().fold(0.0, |acc, &c| acc * x + c)
45}
46
47fn gauss_solve(n: usize, a: &mut [f64], b: &mut [f64]) -> Option<Vec<f64>> {
50 for col in 0..n {
51 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; }
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 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 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 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 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 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}