Skip to main content

qualia_core_db/solvers/learning/splines/
mod.rs

1//! Regression splines & polynomial regression (ISL ch 7) — flexible non-linear
2//! fits expressed as a linear model in a fixed basis, then solved by OLS
3//! (`learning::regression::linear`, no new solver).
4//!
5//! A degree-`d` spline with interior knots `k₁…k_K` uses the **truncated power
6//! basis** `[1, x, …, xᵈ, (x−k₁)ᵈ₊, …, (x−k_K)ᵈ₊]`; polynomial regression is the
7//! special case with no knots. The basis columns form the design matrix; the fit is
8//! ordinary least squares over them (kernel-class `DenseLinear`).
9
10pub 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/// A fitted regression spline (or polynomial, when `knots` is empty).
19#[derive(Debug, Clone)]
20pub struct RegressionSpline {
21    pub degree: usize,
22    pub knots: Vec<f64>,
23    pub coefficients: Vec<f64>,
24}
25
26/// Number of basis columns for `degree` and `n_knots` interior knots.
27pub(crate) fn basis_len(degree: usize, n_knots: usize) -> usize {
28    (degree + 1) + n_knots
29}
30
31/// Evaluate the truncated-power basis row for a scalar `x` into `out`
32/// (length `basis_len`).
33pub(crate) fn basis_row(x: f64, degree: usize, knots: &[f64], out: &mut [f64]) {
34    // Polynomial part 1, x, …, xᵈ.
35    let mut pw = 1.0;
36    for c in out.iter_mut().take(degree + 1) {
37        *c = pw;
38        pw *= x;
39    }
40    // Truncated power terms (x − kⱼ)ᵈ₊.
41    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    /// Fit a degree-`degree` regression spline of `y` on scalar `x` (length `n`)
49    /// with the given interior `knots`. `degree = 3` is the usual cubic spline;
50    /// `knots = []` gives polynomial regression. Fails closed via the OLS solver.
51    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        // Build the n × m design matrix (basis already contains the constant column,
63        // so OLS is fit WITHOUT an extra intercept).
64        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    /// Predict at a scalar `x`.
77    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    /// Predict over a slice of scalar inputs.
85    pub fn predict(&self, x: &[f64]) -> Vec<f64> {
86        x.iter().map(|&xi| self.predict_one(xi)).collect()
87    }
88}
89
90/// Convenience: degree-`degree` polynomial regression (a spline with no knots).
91pub 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        // y = 2 − 3x + x².
108        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        // A curve with a change of behaviour at x≈5 — a cubic spline with a knot
120        // there fits it far better than a single global cubic could.
121        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}