solver/linalg/

regression.rs

use super::gaussian::solve_linear_system;
use faer::prelude::*;

/// Defines a common interface for solving Linear Least Squares problems: $\min_x ||Ax - b||^2$.
///
/// Implementations may vary by:
/// 1. **Memory Usage**: Streaming (low memory) vs Full Matrix (high memory).
/// 2. **Stability**: Normal Equations (unstable) vs QR/SVD (stable).
pub trait RegressionSolver: Sized {
    /// Creates a new solver instance initialized for `n_basis` features.
    fn create(n_basis: usize) -> Self;

    /// Adds a single observation (row of A, element of b) to the solver.
    /// * `features`: The vector of basis function values $\phi(X)$ (row of design matrix).
    /// * `target`: The target value $Y$.
    fn add_observation(&mut self, features: &[f64], target: f64);

    /// Solves the regression problem and returns the coefficients.
    /// Consumes the solver state.
    fn solve(self, lambda: f64) -> Vec<f64>;

    /// Merges another solver state into this one (for parallel reduction).
    fn merge(&mut self, other: Self);
}

/// Normal Equations Solver (Streaming)
///
/// Solves $(A^T A) x = A^T b$.
/// * **Pros**: Low constant memory usage ($O(k^2)$ where $k$ is basis size). Parallelizable via summation.
/// * **Cons**: Numerically unstable if condition number is high ($\kappa(A^T A) = \kappa(A)^2$).
#[derive(Clone, Debug)]
pub struct NormalEquationSolver {
    n_basis: usize,
    /// The matrix $A^T A$ (flattened, size $n \times n$).
    ata: Vec<f64>,
    /// The vector $A^T b$ (size $n$).
    atb: Vec<f64>,
}

impl NormalEquationSolver {
    pub fn new(n_basis: usize) -> Self {
        Self::create(n_basis)
    }
}

impl RegressionSolver for NormalEquationSolver {
    fn create(n_basis: usize) -> Self {
        Self {
            n_basis,
            ata: vec![0.0; n_basis * n_basis],
            atb: vec![0.0; n_basis],
        }
    }

    fn add_observation(&mut self, features: &[f64], target: f64) {
        // Update A^T b: \sum \phi_i * y
        for i in 0..self.n_basis {
            // Using unsafe for performance in hot loop if bounds checks prove costly,
            // but standard indexing is usually fine.
            let feat_i = features[i];
            self.atb[i] += feat_i * target;

            // Update A^T A: \sum \phi_i * \phi_j
            // Exploit symmetry: We only iterate j >= i
            // We store row-major, so index is i * n_basis + j
            let row_offset = i * self.n_basis;

            // Diagonal element (j=i)
            self.ata[row_offset + i] += feat_i * feat_i;

            // Off-diagonal elements
            for (j, &feat_j) in features.iter().enumerate().skip(i + 1) {
                let val = feat_i * feat_j;
                self.ata[row_offset + j] += val;
                self.ata[j * self.n_basis + i] += val; // Symmetric update
            }
        }
    }

    fn solve(mut self, lambda: f64) -> Vec<f64> {
        let n = self.n_basis;
        let mut d = vec![1.0; n];
        for (i, d_i) in d.iter_mut().enumerate() {
            let diag = self.ata[i * n + i];
            if diag > 1e-12 {
                *d_i = 1.0 / diag.sqrt();
            } else {
                *d_i = 1.0; // avoid div by zero, feature is zero anyway
            }
        }

        // Scale ata and atb
        for i in 0..n {
            for j in 0..n {
                self.ata[i * n + j] *= d[i] * d[j];
            }
            self.atb[i] *= d[i];
        }

        // Tikhonov Regularization (Ridge Regression) proportional to unit scale
        if lambda > 0.0 {
            for i in 0..n {
                self.ata[i * n + i] += lambda;
            }
        }

        let x_tilde = solve_linear_system(&self.ata, &self.atb, n);

        let mut x = vec![0.0; n];
        for i in 0..n {
            x[i] = x_tilde[i] * d[i];
        }
        x
    }

    fn merge(&mut self, other: Self) {
        for i in 0..self.ata.len() {
            self.ata[i] += other.ata[i];
        }
        for i in 0..self.atb.len() {
            self.atb[i] += other.atb[i];
        }
    }
}

// Note: A FullMatrixSolver (QR/SVD) is difficult to implement purely based on this trait
// with Rayon's `reduce` because merging two large matrices is expensive (O(N*M)).
// The standard 'reduce' pattern favors the Streaming/NormalEquation approach.
//
// If rigorous stability is required (QR), we generally collect all data first
// and then optimize, rather than streaming into a solver.

/// Full Matrix Solver (Collects all data, then solves via QR Decomposition)
///
/// Solves $\min ||Ax - b||$.
/// * **Pros**: Numerically stable. Correctly handles ill-conditioned matrices.
/// * **Cons**: High memory usage ($O(N \times k)$ where $N$ is number of paths).
#[derive(Clone, Debug)]
pub struct FullMatrixSolver {
    n_basis: usize,
    /// The design matrix $A$ (flattened, row-major: rows=samples, cols=basis).
    a: Vec<f64>,
    /// The target vector $b$ (size=samples).
    b: Vec<f64>,
}

impl FullMatrixSolver {
    pub fn new(n_basis: usize) -> Self {
        Self::create(n_basis)
    }
}

impl RegressionSolver for FullMatrixSolver {
    fn create(n_basis: usize) -> Self {
        Self {
            n_basis,
            a: Vec::new(),
            b: Vec::new(),
        }
    }

    fn add_observation(&mut self, features: &[f64], target: f64) {
        self.a.extend_from_slice(features);
        self.b.push(target);
    }

    fn solve(self, _lambda: f64) -> Vec<f64> {
        let n_samples = self.b.len();
        let n_features = self.n_basis;

        if n_samples < n_features {
            return vec![0.0; n_features];
        }

        // Optimization: Convert to Column-Major layout for cache efficiency.
        // MGS iterates column-by-column. Row-major storage means strided access (bad).
        // Transpose A -> A_col (size: features x samples)
        // This costs one allocation + copy, but speeds up the O(N*K^2) algorithm significantly.

        let mut v_col = vec![0.0; n_samples * n_features];
        // v_col[j * n_samples + i] corresponds to A[i, j]

        for i in 0..n_samples {
            for j in 0..n_features {
                v_col[j * n_samples + i] = self.a[i * n_features + j];
            }
        }

        // Q will also be stored column-major for consistent access
        let mut q_col = vec![0.0; n_samples * n_features];
        let mut r = vec![0.0; n_features * n_features]; // R is small, layout matters less

        for i in 0..n_features {
            let v_i_start = i * n_samples;
            let v_i_end = v_i_start + n_samples;
            let v_i = &v_col[v_i_start..v_i_end];

            // 1. Compute Norm of current column v_i
            // Linear scan!
            let mut norm_sq = 0.0;
            for &val in v_i.iter() {
                norm_sq += val * val;
            }
            let norm = norm_sq.sqrt();
            r[i * n_features + i] = norm;

            // 2. Normalize to get q_i
            let q_i_start = i * n_samples;
            // q_col slice reference (mutable impossible while iterating, index manually)

            if norm > 1e-12 {
                let inv_norm = 1.0 / norm;
                for k in 0..n_samples {
                    q_col[q_i_start + k] = v_col[v_i_start + k] * inv_norm;
                }
            }

            // 3. Orthogonalize remaining columns against q_i
            for j in (i + 1)..n_features {
                let v_j_start = j * n_samples;

                // r_{ij} = <q_i, v_j>
                let mut projection = 0.0;
                let v_j = &v_col[v_j_start..(v_j_start + n_samples)];
                let q_i = &q_col[q_i_start..(q_i_start + n_samples)];

                // Vectorizable loop
                for k in 0..n_samples {
                    projection += q_i[k] * v_j[k];
                }
                r[i * n_features + j] = projection;

                // v_j = v_j - r_{ij} * q_i
                // Vectorizable loop
                let v_j_mut = &mut v_col[v_j_start..(v_j_start + n_samples)];
                for k in 0..n_samples {
                    v_j_mut[k] -= projection * q_i[k];
                }
            }
        }

        // Solve Rx = Q^T b
        // Q is column-major, so Q^T is effectively "row-major" if we view Q as tranposed?
        // Actually, (Q^T b)_i = q_i \cdot b
        // q_i is contiguous in q_col. So this is a dot product. Linear scan.

        let mut y_rhs = vec![0.0; n_features];
        for (i, y_rhs_i) in y_rhs.iter_mut().enumerate() {
            let q_i_start = i * n_samples;
            let mut sum = 0.0;
            for (&a, &b) in q_col[q_i_start..].iter().zip(self.b.iter()) {
                sum += a * b;
            }
            *y_rhs_i = sum;
        }

        // 2. Solve Rx = y by Back Substitution
        let mut x = vec![0.0; n_features];
        for i in (0..n_features).rev() {
            let mut sum = y_rhs[i];
            for j in (i + 1)..n_features {
                sum -= r[i * n_features + j] * x[j];
            }
            if r[i * n_features + i].abs() > 1e-12 {
                x[i] = sum / r[i * n_features + i];
            } else {
                x[i] = 0.0;
            }
        }

        x
    }

    fn merge(&mut self, mut other: Self) {
        self.a.append(&mut other.a);
        self.b.append(&mut other.b);
    }
}

/// Faer Solver (Pure Rust High-Performance)
///
/// Uses `faer` crate for QR decomposition.
#[derive(Clone, Debug)]
pub struct FaerSolver {
    n_basis: usize,
    a: Vec<f64>,
    b: Vec<f64>,
}

impl RegressionSolver for FaerSolver {
    fn create(n_basis: usize) -> Self {
        Self {
            n_basis,
            a: Vec::new(),
            b: Vec::new(),
        }
    }

    fn add_observation(&mut self, features: &[f64], target: f64) {
        self.a.extend_from_slice(features);
        self.b.push(target);
    }

    fn solve(self, lambda: f64) -> Vec<f64> {
        let n_samples = self.b.len();
        let n_features = self.n_basis;
        if n_samples < n_features {
            return vec![0.0; n_features];
        }

        // faer generally prefers column-major matrices.
        // 1. Construct Matrix A (col-major)
        let mut a_col = vec![0.0; n_samples * n_features];
        for i in 0..n_samples {
            for j in 0..n_features {
                a_col[j * n_samples + i] = self.a[i * n_features + j];
            }
        }

        // Tikhonov regularization: augment A with sqrt(lambda)*I and b with zeros
        // This transforms min||Ax-b||^2 + lambda*||x||^2 into min||[A; sqrt(lambda)*I]x - [b; 0]||^2
        if lambda > 0.0 {
            let sqrt_lam = lambda.sqrt();
            let total_rows = n_samples + n_features;
            let mut a_aug = vec![0.0; total_rows * n_features];
            // Copy original A columns
            for j in 0..n_features {
                for i in 0..n_samples {
                    a_aug[j * total_rows + i] = a_col[j * n_samples + i];
                }
                // Add sqrt(lambda) on diagonal of augmented part
                a_aug[j * total_rows + n_samples + j] = sqrt_lam;
            }
            let mut b_aug = Vec::with_capacity(total_rows);
            b_aug.extend_from_slice(&self.b);
            b_aug.resize(total_rows, 0.0);

            let a_mat = faer::mat::from_column_major_slice::<f64>(&a_aug, total_rows, n_features);
            let b_mat = faer::mat::from_column_major_slice::<f64>(&b_aug, total_rows, 1);
            let qr = a_mat.col_piv_qr();
            let x = qr.solve_lstsq(b_mat);
            return (0..n_features).map(|i| x.read(i, 0)).collect();
        }

        // No regularization path
        let a_mat = faer::mat::from_column_major_slice::<f64>(&a_col, n_samples, n_features);
        let b_mat = faer::mat::from_column_major_slice::<f64>(&self.b, n_samples, 1);
        let qr = a_mat.col_piv_qr();
        let x = qr.solve_lstsq(b_mat);
        (0..n_features).map(|i| x.read(i, 0)).collect()
    }

    fn merge(&mut self, mut other: Self) {
        self.a.append(&mut other.a);
        self.b.append(&mut other.b);
    }
}