qualia_core_db/solvers/learning/splines/
gam.rs

1//! Generalized Additive Model (ISL ch 7.7) — `y = β₀ + Σⱼ fⱼ(xⱼ)`, each `fⱼ` a
2//! regression spline, fit by **backfitting**: cycle through the features, fitting
3//! each smooth term to the partial residual of all the others. Reuses
4//! [`super::RegressionSpline`] (no duplicated basis/OLS). Kernel-class `DenseLinear`.
5
6use crate::solvers::learning::splines::RegressionSpline;
7use crate::solvers::learning::LearningError;
8use crate::solvers::statistics::descriptive::mean;
9
10/// A fitted additive model: an intercept plus one centered smooth term per feature.
11#[derive(Debug, Clone)]
12pub struct Gam {
13    pub intercept: f64,
14    /// One `(spline, mean_offset)` per feature; the term value is
15    /// `spline(x) − mean_offset` (centered for identifiability).
16    terms: Vec<(RegressionSpline, f64)>,
17    p: usize,
18}
19
20impl Gam {
21    /// Fit by backfitting. `degree` and `knots_per_feature[j]` define each feature's
22    /// spline; `max_iter` backfitting sweeps. Fails closed on shape mismatch.
23    pub fn fit(
24        x: &[f64],
25        y: &[f64],
26        n: usize,
27        p: usize,
28        degree: usize,
29        knots_per_feature: &[Vec<f64>],
30        max_iter: usize,
31    ) -> Result<Self, LearningError> {
32        if n == 0 || p == 0 || x.len() != n * p || y.len() != n || knots_per_feature.len() != p {
33            return Err(LearningError::InvalidDimension);
34        }
35        let intercept = mean(y).ok_or(LearningError::InsufficientData)?;
36
37        // Per-feature columns (contiguous) for spline fitting.
38        let cols: Vec<Vec<f64>> = (0..p)
39            .map(|j| (0..n).map(|i| x[i * p + j]).collect::<Vec<f64>>())
40            .collect();
41
42        // Current fitted-term values at each training point, n×p (start at 0).
43        let mut term_vals = vec![0.0; n * p];
44        // Placeholder splines (degree-`degree`, replaced during the first sweep).
45        let mut terms: Vec<(RegressionSpline, f64)> = Vec::with_capacity(p);
46        for j in 0..p {
47            let spline =
48                RegressionSpline::fit(&cols[j], &vec![0.0; n], n, degree, &knots_per_feature[j])?;
49            terms.push((spline, 0.0));
50        }
51
52        for _ in 0..max_iter.max(1) {
53            for j in 0..p {
54                // Partial residual: y − intercept − Σ_{k≠j} f_k.
55                let mut resid = vec![0.0; n];
56                for i in 0..n {
57                    let mut s = y[i] - intercept;
58                    for k in 0..p {
59                        if k != j {
60                            s -= term_vals[i * p + k];
61                        }
62                    }
63                    resid[i] = s;
64                }
65                // Fit f_j to the residual, then center it (mean 0 over the data).
66                let spline =
67                    RegressionSpline::fit(&cols[j], &resid, n, degree, &knots_per_feature[j])?;
68                let raw: Vec<f64> = cols[j].iter().map(|&xi| spline.predict_one(xi)).collect();
69                let offset = mean(&raw).unwrap_or(0.0);
70                for i in 0..n {
71                    term_vals[i * p + j] = raw[i] - offset;
72                }
73                terms[j] = (spline, offset);
74            }
75        }
76
77        Ok(Self {
78            intercept,
79            terms,
80            p,
81        })
82    }
83
84    /// Predict for one feature row.
85    pub fn predict_row(&self, x_row: &[f64]) -> f64 {
86        let mut s = self.intercept;
87        for (j, (spline, offset)) in self.terms.iter().enumerate() {
88            s += spline.predict_one(x_row[j]) - offset;
89        }
90        s
91    }
92
93    pub fn predict(&self, x: &[f64], n: usize) -> Vec<f64> {
94        (0..n)
95            .map(|i| self.predict_row(&x[i * self.p..(i + 1) * self.p]))
96            .collect()
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::solvers::learning::metrics::regression::r2_score;
104
105    #[test]
106    fn fits_an_additive_nonlinear_surface() {
107        // y = sin(x0) + 0.1·x1²  — additive but nonlinear in each feature.
108        let n = 40;
109        let mut x = vec![0.0; n * 2];
110        let mut y = vec![0.0; n];
111        for i in 0..n {
112            let x0 = (i as f64) * 0.15;
113            let x1 = ((i * 7) % 40) as f64 * 0.1;
114            x[i * 2] = x0;
115            x[i * 2 + 1] = x1;
116            y[i] = x0.sin() + 0.1 * x1 * x1;
117        }
118        let knots = vec![vec![1.5, 3.0, 4.5], vec![1.0, 2.0, 3.0]];
119        let gam = Gam::fit(&x, &y, n, 2, 3, &knots, 10).unwrap();
120        let preds = gam.predict(&x, n);
121        assert!(
122            r2_score(&y, &preds).unwrap() > 0.95,
123            "GAM should fit the additive surface"
124        );
125    }
126
127    #[test]
128    fn recovers_a_linear_additive_model() {
129        // y = 2·x0 − 3·x1 + 1: an additive (linear) model; GAM recovers it.
130        let n = 25;
131        let mut x = vec![0.0; n * 2];
132        let mut y = vec![0.0; n];
133        for i in 0..n {
134            let x0 = i as f64;
135            let x1 = (i % 5) as f64;
136            x[i * 2] = x0;
137            x[i * 2 + 1] = x1;
138            y[i] = 2.0 * x0 - 3.0 * x1 + 1.0;
139        }
140        let knots = vec![vec![], vec![]]; // no knots → linear terms
141        let gam = Gam::fit(&x, &y, n, 2, 1, &knots, 20).unwrap();
142        let preds = gam.predict(&x, n);
143        assert!(r2_score(&y, &preds).unwrap() > 0.999);
144    }
145
146    #[test]
147    fn guards() {
148        assert_eq!(
149            Gam::fit(&[1.0, 2.0], &[1.0], 2, 1, 3, &[vec![]], 5).unwrap_err(),
150            LearningError::InvalidDimension
151        );
152    }
153}
qualia_core_db/solvers/learning/splines/gam.rs

qualia_core_db/solvers/learning/splines/
gam.rs
