Skip to main content

qualia_core_db/solvers/learning/regression/
pls.rs

1//! Partial Least Squares regression (ISL ch 6.3.2) — PLS1 (univariate response) by
2//! the NIPALS algorithm. Unlike PCR (which picks directions of high predictor
3//! variance), PLS picks directions of high *covariance with the response*. The small
4//! component-space solve reuses `linear_algebra::qr` (no new solver). Kernel-class
5//! `DenseLinear`.
6
7use crate::solvers::learning::LearningError;
8use crate::solvers::linear_algebra::qr::{qr_factor, qr_solve_least_squares};
9use crate::solvers::statistics::descriptive::mean;
10
11/// A fitted PLS model collapsed to predictor-space coefficients + intercept.
12#[derive(Debug, Clone)]
13pub struct PlsModel {
14    pub coefficients: Vec<f64>,
15    pub intercept: f64,
16    pub n_components: usize,
17}
18
19impl PlsModel {
20    pub fn predict_row(&self, x_row: &[f64]) -> f64 {
21        self.intercept
22            + self
23                .coefficients
24                .iter()
25                .zip(x_row)
26                .map(|(b, x)| b * x)
27                .sum::<f64>()
28    }
29    pub fn predict(&self, x: &[f64], n: usize, p: usize) -> Vec<f64> {
30        (0..n)
31            .map(|i| self.predict_row(&x[i * p..(i + 1) * p]))
32            .collect()
33    }
34}
35
36/// Fit PLS1 with `n_components` (clamped to `p`). Fails closed on shape mismatch /
37/// `n < 2` / a degenerate (zero-covariance) extraction.
38pub fn fit(
39    x: &[f64],
40    y: &[f64],
41    n: usize,
42    p: usize,
43    n_components: usize,
44) -> Result<PlsModel, LearningError> {
45    if n == 0 || p == 0 || x.len() != n * p || y.len() != n {
46        return Err(LearningError::InvalidDimension);
47    }
48    if n < 2 {
49        return Err(LearningError::InsufficientData);
50    }
51    let a = n_components.clamp(1, p);
52
53    // Centre X and y.
54    let mut xbar = vec![0.0; p];
55    let mut colbuf = vec![0.0; n];
56    for j in 0..p {
57        for i in 0..n {
58            colbuf[i] = x[i * p + j];
59        }
60        xbar[j] = mean(&colbuf).ok_or(LearningError::InsufficientData)?;
61    }
62    let ybar = mean(y).ok_or(LearningError::InsufficientData)?;
63    let mut xc = vec![0.0; n * p];
64    let mut yc = vec![0.0; n];
65    for i in 0..n {
66        for j in 0..p {
67            xc[i * p + j] = x[i * p + j] - xbar[j];
68        }
69        yc[i] = y[i] - ybar;
70    }
71
72    // NIPALS: collect weight (w), loading (p_load) columns and scalar q per component.
73    let mut w_mat = vec![0.0; p * a]; // p × a, column per component
74    let mut p_mat = vec![0.0; p * a];
75    let mut q_vec = vec![0.0; a];
76
77    for comp in 0..a {
78        // w = Xcᵀ yc, normalized.
79        let mut w = vec![0.0; p];
80        for j in 0..p {
81            let mut s = 0.0;
82            for i in 0..n {
83                s += xc[i * p + j] * yc[i];
84            }
85            w[j] = s;
86        }
87        let wn = (w.iter().map(|v| v * v).sum::<f64>()).sqrt();
88        if wn < 1e-12 {
89            // No remaining covariance — stop early (use the components found so far).
90            if comp == 0 {
91                return Err(LearningError::InsufficientData);
92            }
93            return finalize(&w_mat, &p_mat, &q_vec, comp, p, &xbar, ybar);
94        }
95        for v in w.iter_mut() {
96            *v /= wn;
97        }
98        // Scores t = Xc w.
99        let mut t = vec![0.0; n];
100        for i in 0..n {
101            t[i] = (0..p).map(|j| xc[i * p + j] * w[j]).sum();
102        }
103        let tt: f64 = t.iter().map(|v| v * v).sum();
104        if tt < 1e-12 {
105            return finalize(&w_mat, &p_mat, &q_vec, comp, p, &xbar, ybar);
106        }
107        // Loadings p_load = Xcᵀ t / tt, q = ycᵀ t / tt.
108        let mut p_load = vec![0.0; p];
109        for j in 0..p {
110            let mut s = 0.0;
111            for i in 0..n {
112                s += xc[i * p + j] * t[i];
113            }
114            p_load[j] = s / tt;
115        }
116        let q = yc.iter().zip(&t).map(|(yi, ti)| yi * ti).sum::<f64>() / tt;
117        // Deflate.
118        for i in 0..n {
119            for j in 0..p {
120                xc[i * p + j] -= t[i] * p_load[j];
121            }
122            yc[i] -= t[i] * q;
123        }
124        for j in 0..p {
125            w_mat[j * a + comp] = w[j];
126            p_mat[j * a + comp] = p_load[j];
127        }
128        q_vec[comp] = q;
129    }
130
131    finalize(&w_mat, &p_mat, &q_vec, a, p, &xbar, ybar)
132}
133
134/// Collapse `a` components into predictor-space coefficients
135/// `β = W (PᵀW)⁻¹ q`, intercept `ȳ − x̄ᵀβ`.
136fn finalize(
137    w_mat: &[f64],
138    p_mat: &[f64],
139    q_vec: &[f64],
140    a: usize,
141    p: usize,
142    xbar: &[f64],
143    ybar: f64,
144) -> Result<PlsModel, LearningError> {
145    // M = PᵀW (a×a): M[r,c] = Σ_j p_mat[j,r]·w_mat[j,c].
146    let stride = w_mat.len() / p; // = number of allocated components (>= a)
147    let mut m = vec![0.0; a * a];
148    for r in 0..a {
149        for c in 0..a {
150            let mut s = 0.0;
151            for j in 0..p {
152                s += p_mat[j * stride + r] * w_mat[j * stride + c];
153            }
154            m[r * a + c] = s;
155        }
156    }
157    // Solve M α = q (square) via QR.
158    let mut tau = vec![0.0; a];
159    qr_factor(a, a, &mut m, &mut tau)?;
160    let mut b = q_vec[..a].to_vec();
161    let mut alpha = vec![0.0; a];
162    qr_solve_least_squares(a, a, &m, &tau, &mut b, &mut alpha)?;
163    // β = W α.
164    let mut coefficients = vec![0.0; p];
165    for j in 0..p {
166        let mut s = 0.0;
167        for c in 0..a {
168            s += w_mat[j * stride + c] * alpha[c];
169        }
170        coefficients[j] = s;
171    }
172    let intercept = ybar
173        - coefficients
174            .iter()
175            .zip(xbar)
176            .map(|(b, m)| b * m)
177            .sum::<f64>();
178    Ok(PlsModel {
179        coefficients,
180        intercept,
181        n_components: a,
182    })
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use crate::solvers::learning::regression::linear;
189
190    #[test]
191    fn full_components_matches_ols() {
192        // With min(n-1, p) components PLS reproduces the OLS fit.
193        let x = [1.0, 2.0, 2.0, 1.0, 3.0, 0.0, 4.0, 5.0, 5.0, 4.0, 6.0, 1.0];
194        let y = [3.0, 5.0, 4.0, 9.0, 13.0, 8.0];
195        let pls = fit(&x, &y, 6, 2, 2).unwrap();
196        let ols = linear::fit(&x, &y, 6, 2, true).unwrap();
197        let p_pred = pls.predict(&x, 6, 2);
198        let o_pred = ols.predict(&x, 6, 2);
199        for i in 0..6 {
200            assert!(
201                (p_pred[i] - o_pred[i]).abs() < 1e-6,
202                "{} vs {}",
203                p_pred[i],
204                o_pred[i]
205            );
206        }
207    }
208
209    #[test]
210    fn one_component_tracks_covariance_direction() {
211        // y driven by x0; PLS with 1 component fits well (it targets covariance).
212        let n = 20;
213        let mut x = vec![0.0; n * 2];
214        let mut y = vec![0.0; n];
215        for i in 0..n {
216            let t = i as f64;
217            x[i * 2] = t;
218            x[i * 2 + 1] = ((i % 5) as f64) - 2.0; // unrelated to y
219            y[i] = 3.0 * t + 2.0;
220        }
221        let pls = fit(&x, &y, n, 2, 1).unwrap();
222        let preds = pls.predict(&x, n, 2);
223        let r2 = crate::solvers::learning::metrics::regression::r2_score(&y, &preds).unwrap();
224        assert!(r2 > 0.99, "r2 {r2}");
225    }
226
227    #[test]
228    fn guards() {
229        assert_eq!(
230            fit(&[1.0, 2.0, 3.0], &[1.0, 2.0], 2, 2, 1).unwrap_err(),
231            LearningError::InvalidDimension
232        );
233    }
234}