qualia_core_db/solvers/learning/splines/
mod.rs1pub mod gam;
11pub mod smoothing;
12pub use gam::Gam;
13pub use smoothing::SmoothingSpline;
14
15use crate::solvers::learning::regression::linear;
16use crate::solvers::learning::LearningError;
17
18#[derive(Debug, Clone)]
20pub struct RegressionSpline {
21 pub degree: usize,
22 pub knots: Vec<f64>,
23 pub coefficients: Vec<f64>,
24}
25
26pub(crate) fn basis_len(degree: usize, n_knots: usize) -> usize {
28 (degree + 1) + n_knots
29}
30
31pub(crate) fn basis_row(x: f64, degree: usize, knots: &[f64], out: &mut [f64]) {
34 let mut pw = 1.0;
36 for c in out.iter_mut().take(degree + 1) {
37 *c = pw;
38 pw *= x;
39 }
40 for (j, &k) in knots.iter().enumerate() {
42 let d = x - k;
43 out[degree + 1 + j] = if d > 0.0 { d.powi(degree as i32) } else { 0.0 };
44 }
45}
46
47impl RegressionSpline {
48 pub fn fit(
52 x: &[f64],
53 y: &[f64],
54 n: usize,
55 degree: usize,
56 knots: &[f64],
57 ) -> Result<Self, LearningError> {
58 if n == 0 || x.len() != n || y.len() != n || degree == 0 {
59 return Err(LearningError::InvalidDimension);
60 }
61 let m = basis_len(degree, knots.len());
62 let mut design = vec![0.0; n * m];
65 for i in 0..n {
66 basis_row(x[i], degree, knots, &mut design[i * m..(i + 1) * m]);
67 }
68 let model = linear::fit(&design, y, n, m, false)?;
69 Ok(Self {
70 degree,
71 knots: knots.to_vec(),
72 coefficients: model.coefficients,
73 })
74 }
75
76 pub fn predict_one(&self, x: f64) -> f64 {
78 let m = basis_len(self.degree, self.knots.len());
79 let mut row = vec![0.0; m];
80 basis_row(x, self.degree, &self.knots, &mut row);
81 row.iter().zip(&self.coefficients).map(|(b, c)| b * c).sum()
82 }
83
84 pub fn predict(&self, x: &[f64]) -> Vec<f64> {
86 x.iter().map(|&xi| self.predict_one(xi)).collect()
87 }
88}
89
90pub fn polynomial_regression(
92 x: &[f64],
93 y: &[f64],
94 n: usize,
95 degree: usize,
96) -> Result<RegressionSpline, LearningError> {
97 RegressionSpline::fit(x, y, n, degree, &[])
98}
99
100#[cfg(test)]
101mod tests {
102 use super::*;
103 use crate::solvers::learning::metrics::regression::r2_score;
104
105 #[test]
106 fn polynomial_recovers_a_quadratic_exactly() {
107 let x: Vec<f64> = (0..8).map(|i| i as f64).collect();
109 let y: Vec<f64> = x.iter().map(|&xi| 2.0 - 3.0 * xi + xi * xi).collect();
110 let m = polynomial_regression(&x, &y, 8, 2).unwrap();
111 assert!((m.coefficients[0] - 2.0).abs() < 1e-6);
112 assert!((m.coefficients[1] + 3.0).abs() < 1e-6);
113 assert!((m.coefficients[2] - 1.0).abs() < 1e-6);
114 assert!((m.predict_one(10.0) - (2.0 - 30.0 + 100.0)).abs() < 1e-5);
115 }
116
117 #[test]
118 fn cubic_spline_fits_a_kinked_curve() {
119 let n = 30;
122 let x: Vec<f64> = (0..n).map(|i| i as f64 / 3.0).collect();
123 let y: Vec<f64> = x
124 .iter()
125 .map(|&xi| {
126 if xi < 5.0 {
127 (xi).sin()
128 } else {
129 0.3 * (xi - 5.0) + (5.0f64).sin()
130 }
131 })
132 .collect();
133 let m = RegressionSpline::fit(&x, &y, n, 3, &[3.0, 5.0, 7.0]).unwrap();
134 let preds = m.predict(&x);
135 assert!(
136 r2_score(&y, &preds).unwrap() > 0.97,
137 "spline should fit well"
138 );
139 }
140
141 #[test]
142 fn guards() {
143 assert_eq!(
144 RegressionSpline::fit(&[1.0, 2.0], &[1.0], 2, 3, &[]).unwrap_err(),
145 LearningError::InvalidDimension
146 );
147 assert_eq!(
148 RegressionSpline::fit(&[1.0], &[1.0], 1, 0, &[]).unwrap_err(),
149 LearningError::InvalidDimension
150 );
151 }
152}