qualia_core_db/solvers/learning/regression/
pls.rs1use crate::solvers::learning::LearningError;
8use crate::solvers::linear_algebra::qr::{qr_factor, qr_solve_least_squares};
9use crate::solvers::statistics::descriptive::mean;
10
11#[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
36pub 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 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 let mut w_mat = vec![0.0; p * a]; 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 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 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 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 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 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
134fn 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 let stride = w_mat.len() / p; 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 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 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 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 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; 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}