qualia_core_db/solvers/learning/regression/
linear.rs

1//! Multiple linear regression (ISL ch 3) — ordinary least squares with full
2//! inference, solved through the engine's linear-algebra library (no re-implemented
3//! solver) and the statistics distributions (real p-values).
4//!
5//! Fit `y = β₀ + β₁x₁ + … + β_p x_p` by the normal equations `(XᵀX)β = Xᵀy`, formed
6//! with `linear_algebra::gemm`/`matvec` and solved (and inverted, for the coefficient
7//! standard errors) with `linear_algebra::cholesky`. Inference — t-tests on each
8//! coefficient and the overall F-test — uses `statistics::distributions`.
9//!
10//! Kernel-class: `DenseLinear` (the GEMM/solve), so it is dispatch-ready against
11//! `ComputePolicy`; for the small p×p normal-equations solve the CPU path is the
12//! right one — the GPU win is in forming `XᵀX` for large n, wired with the bridge.
13
14use crate::solvers::learning::LearningError;
15use crate::solvers::linear_algebra::cholesky::{cholesky_factor, cholesky_solve};
16use crate::solvers::linear_algebra::gemm::{gemm, matvec, Transpose};
17use crate::solvers::statistics::descriptive::mean;
18use crate::solvers::statistics::distributions::{fisher_f, students_t};
19
20/// A fitted OLS model with inferential output. When `fit_intercept` is true,
21/// `coefficients[0]` is the intercept and `coefficients[1..]` align with the
22/// predictor columns; the `*_per_coef` vectors are aligned the same way.
23#[derive(Debug, Clone)]
24pub struct LinearModel {
25    pub coefficients: Vec<f64>,
26    pub fit_intercept: bool,
27    pub std_errors: Vec<f64>,
28    pub t_values: Vec<f64>,
29    pub p_values: Vec<f64>,
30    pub r_squared: f64,
31    pub adj_r_squared: f64,
32    /// Overall F-statistic (all slopes = 0) and its p-value. `None` without an intercept.
33    pub f_statistic: Option<f64>,
34    pub f_p_value: Option<f64>,
35    pub residual_std_error: f64,
36    pub df_residual: usize,
37    pub n: usize,
38}
39
40impl LinearModel {
41    /// Predict for one feature row (length `p`, predictors only — the intercept is
42    /// applied internally).
43    pub fn predict_row(&self, x_row: &[f64]) -> f64 {
44        let (b0, betas) = if self.fit_intercept {
45            (self.coefficients[0], &self.coefficients[1..])
46        } else {
47            (0.0, &self.coefficients[..])
48        };
49        b0 + betas.iter().zip(x_row).map(|(b, x)| b * x).sum::<f64>()
50    }
51
52    /// Predict for a row-major `n × p` feature matrix.
53    pub fn predict(&self, x: &[f64], n: usize, p: usize) -> Vec<f64> {
54        (0..n)
55            .map(|i| self.predict_row(&x[i * p..(i + 1) * p]))
56            .collect()
57    }
58}
59
60/// Fit OLS of `y` (length `n`) on a row-major `n × p` predictor matrix `x`.
61/// `fit_intercept` prepends a constant column. Fails closed:
62/// `InvalidDimension` on a shape mismatch, `InsufficientData` if `n ≤ params`,
63/// `Singular` on collinear predictors.
64pub fn fit(
65    x: &[f64],
66    y: &[f64],
67    n: usize,
68    p: usize,
69    fit_intercept: bool,
70) -> Result<LinearModel, LearningError> {
71    if n == 0 || p == 0 || x.len() != n * p || y.len() != n {
72        return Err(LearningError::InvalidDimension);
73    }
74    let k = p + usize::from(fit_intercept); // total parameters
75    if n <= k {
76        return Err(LearningError::InsufficientData);
77    }
78
79    // Build the design matrix D (n × k), row-major, with a leading 1s column if
80    // an intercept is fit.
81    let mut d = vec![0.0; n * k];
82    for i in 0..n {
83        let base = i * k;
84        if fit_intercept {
85            d[base] = 1.0;
86            d[base + 1..base + k].copy_from_slice(&x[i * p..(i + 1) * p]);
87        } else {
88            d[base..base + k].copy_from_slice(&x[i * p..(i + 1) * p]);
89        }
90    }
91
92    // Normal equations: A = DᵀD (k×k), b = Dᵀy (k).
93    let mut a = vec![0.0; k * k];
94    gemm(
95        Transpose::Yes,
96        Transpose::No,
97        k,
98        k,
99        n,
100        1.0,
101        &d,
102        &d,
103        0.0,
104        &mut a,
105    )?;
106    let mut b = vec![0.0; k];
107    matvec(Transpose::Yes, k, n, &d, y, &mut b)?;
108
109    // Cholesky factor of the SPD Gram matrix; fail closed if not positive-definite
110    // (collinear predictors).
111    let mut l = vec![0.0; k * k];
112    cholesky_factor(k, &a, &mut l).map_err(|_| LearningError::Singular)?;
113
114    // Coefficients: solve A·β = b.
115    let mut coefficients = vec![0.0; k];
116    cholesky_solve(k, &l, &b, &mut coefficients)?;
117
118    // Residuals and sums of squares.
119    let mut yhat = vec![0.0; n];
120    matvec(Transpose::No, n, k, &d, &coefficients, &mut yhat)?;
121    let ybar = mean(y).ok_or(LearningError::InsufficientData)?;
122    let mut sse = 0.0;
123    let mut sst = 0.0;
124    for i in 0..n {
125        sse += (y[i] - yhat[i]).powi(2);
126        sst += (y[i] - ybar).powi(2);
127    }
128    let df_residual = n - k;
129    let sigma2 = sse / df_residual as f64;
130    let residual_std_error = sigma2.sqrt();
131
132    // (XᵀX)⁻¹ diagonal for coefficient standard errors: solve A·cⱼ = eⱼ via the
133    // existing Cholesky factor and read cⱼ[j].
134    let mut std_errors = vec![0.0; k];
135    let mut t_values = vec![0.0; k];
136    let mut p_values = vec![0.0; k];
137    let df = df_residual as f64;
138    let mut ej = vec![0.0; k];
139    let mut cj = vec![0.0; k];
140    for j in 0..k {
141        ej.iter_mut().for_each(|v| *v = 0.0);
142        ej[j] = 1.0;
143        cholesky_solve(k, &l, &ej, &mut cj)?;
144        let var = sigma2 * cj[j];
145        let se = if var > 0.0 { var.sqrt() } else { 0.0 };
146        std_errors[j] = se;
147        if se > 0.0 {
148            let t = coefficients[j] / se;
149            t_values[j] = t;
150            p_values[j] = students_t::two_sided_p(t, df);
151        } else {
152            t_values[j] = if coefficients[j] == 0.0 {
153                0.0
154            } else {
155                f64::INFINITY
156            };
157            p_values[j] = if coefficients[j] == 0.0 { 1.0 } else { 0.0 };
158        }
159    }
160
161    let r_squared = if sst > 0.0 { 1.0 - sse / sst } else { 1.0 };
162    let adj_r_squared = if df_residual > 0 && sst > 0.0 {
163        1.0 - (1.0 - r_squared) * (n as f64 - 1.0) / df as f64
164    } else {
165        r_squared
166    };
167
168    // Overall F-test (only meaningful with an intercept): F = MSR/MSE.
169    let (f_statistic, f_p_value) = if fit_intercept && p >= 1 && sst > 0.0 {
170        let df_model = p as f64;
171        let f = ((sst - sse) / df_model) / sigma2;
172        (Some(f), Some(fisher_f::upper_p(f, df_model, df)))
173    } else {
174        (None, None)
175    };
176
177    Ok(LinearModel {
178        coefficients,
179        fit_intercept,
180        std_errors,
181        t_values,
182        p_values,
183        r_squared,
184        adj_r_squared,
185        f_statistic,
186        f_p_value,
187        residual_std_error,
188        df_residual,
189        n,
190    })
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn recovers_exact_plane() {
199        // y = 1 + 2·x1 + 3·x2 exactly, 5 points.
200        let x = [0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 1.0];
201        let y = [1.0, 3.0, 4.0, 6.0, 8.0]; // 1+2x1+3x2
202        let m = fit(&x, &y, 5, 2, true).unwrap();
203        assert!(
204            (m.coefficients[0] - 1.0).abs() < 1e-9,
205            "intercept {}",
206            m.coefficients[0]
207        );
208        assert!(
209            (m.coefficients[1] - 2.0).abs() < 1e-9,
210            "b1 {}",
211            m.coefficients[1]
212        );
213        assert!(
214            (m.coefficients[2] - 3.0).abs() < 1e-9,
215            "b2 {}",
216            m.coefficients[2]
217        );
218        assert!((m.r_squared - 1.0).abs() < 1e-12);
219        // Prediction matches.
220        assert!((m.predict_row(&[3.0, 2.0]) - (1.0 + 6.0 + 6.0)).abs() < 1e-9);
221    }
222
223    #[test]
224    fn matches_simple_regression_for_one_predictor() {
225        // Compare against the closed-form simple OLS already in statistics.
226        let x = [1.0, 2.0, 3.0, 4.0, 5.0];
227        let y = [2.1, 3.9, 6.1, 7.9, 10.2];
228        let m = fit(&x, &y, 5, 1, true).unwrap();
229        let simple =
230            crate::solvers::statistics::regression::simple_linear_regression(&x, &y).unwrap();
231        assert!((m.coefficients[0] - simple.intercept).abs() < 1e-9);
232        assert!((m.coefficients[1] - simple.slope).abs() < 1e-9);
233        // Slope p-value agrees with the simple-regression module.
234        assert!((m.p_values[1] - simple.slope_p_value).abs() < 1e-9);
235    }
236
237    #[test]
238    fn detects_collinear_predictors() {
239        // x2 = 2·x1 → singular normal equations → fail closed (no bogus fit).
240        let x = [1.0, 2.0, 2.0, 4.0, 3.0, 6.0, 4.0, 8.0, 5.0, 10.0];
241        let y = [1.0, 2.0, 3.0, 4.0, 5.0];
242        assert_eq!(
243            fit(&x, &y, 5, 2, true).unwrap_err(),
244            LearningError::Singular
245        );
246    }
247
248    #[test]
249    fn guards_insufficient_data() {
250        // 2 samples, 2 predictors + intercept = 3 params → n ≤ k.
251        let x = [1.0, 2.0, 3.0, 4.0];
252        let y = [1.0, 2.0];
253        assert_eq!(
254            fit(&x, &y, 2, 2, true).unwrap_err(),
255            LearningError::InsufficientData
256        );
257    }
258
259    #[test]
260    fn significant_predictor_has_small_p() {
261        // Strong linear signal in x1, noise-free → tiny p-value for its slope.
262        let x = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
263        let y = [2.0, 4.1, 5.9, 8.0, 10.1, 12.0];
264        let m = fit(&x, &y, 6, 1, true).unwrap();
265        assert!(m.p_values[1] < 1e-4);
266        assert!(m.f_p_value.unwrap() < 1e-4);
267    }
268}
qualia_core_db/solvers/learning/regression/linear.rs

qualia_core_db/solvers/learning/regression/
linear.rs
