linreg_core/regularized/

ridge.rs

//! Ridge regression (L2-regularized linear regression).
//!
//! This module provides ridge regression implementation using the augmented QR
//! approach, which is numerically stable and avoids forming X^T X explicitly.
//!
//! # Ridge Regression Objective
//!
//! Ridge regression solves:
//!
//! ```text
//! minimize over (β₀, β):
//!
//!     (1/(2n)) * Σᵢ (yᵢ - β₀ - xᵢᵀβ)² + (λ/2) * ||β||₂²
//! ```
//!
//! The intercept `β₀` is **not penalized**.
//!
//! # Solution Method
//!
//! We use the augmented least-squares approach:
//!
//! ```text
//! minimize || [y; 0] - [X; √λ*I] * β ||²
//! ```
//!
//! This transforms the ridge problem into a standard least squares problem
//! that can be solved with QR decomposition.

use crate::error::{Error, Result};
use crate::linalg::Matrix;
use crate::regularized::preprocess::{
    predict, standardize_xy, unstandardize_coefficients, StandardizeOptions,
};

#[cfg(feature = "wasm")]
use serde::Serialize;

/// Options for ridge regression fitting.
///
/// # Fields
///
/// * `lambda` - Regularization strength (single value)
/// * `intercept` - Whether to include an intercept term (default: true)
/// * `standardize` - Whether to standardize predictors (default: true)
#[derive(Clone, Debug)]
pub struct RidgeFitOptions {
    /// Regularization strength (must be >= 0)
    pub lambda: f64,
    /// Whether to include an intercept
    pub intercept: bool,
    /// Whether to standardize predictors
    pub standardize: bool,
}

impl Default for RidgeFitOptions {
    fn default() -> Self {
        RidgeFitOptions {
            lambda: 1.0,
            intercept: true,
            standardize: true,
        }
    }
}

/// Result of a ridge regression fit.
///
/// # Fields
///
/// * `lambda` - The lambda value used for fitting
/// * `intercept` - Intercept coefficient (on original scale)
/// * `coefficients` - Slope coefficients (on original scale)
/// * `fitted_values` - In-sample predictions
/// * `residuals` - Residuals (y - fitted_values)
/// * `df` - Effective degrees of freedom (trace of H = X(X'X + λI)^(-1)X')
/// * `r_squared` - R² (coefficient of determination)
/// * `adj_r_squared` - Adjusted R² (using effective df)
/// * `mse` - Mean squared error
/// * `rmse` - Root mean squared error
/// * `mae` - Mean absolute error
/// * `standardization_info` - Information about standardization applied
#[derive(Clone, Debug)]
#[cfg_attr(feature = "wasm", derive(Serialize))]
pub struct RidgeFit {
    /// Lambda value used for fitting
    pub lambda: f64,
    /// Intercept on original scale
    pub intercept: f64,
    /// Slope coefficients on original scale
    pub coefficients: Vec<f64>,
    /// Fitted values
    pub fitted_values: Vec<f64>,
    /// Residuals
    pub residuals: Vec<f64>,
    /// Effective degrees of freedom
    pub df: f64,
    /// R² (coefficient of determination)
    pub r_squared: f64,
    /// Adjusted R² (penalized for effective df)
    pub adj_r_squared: f64,
    /// Mean squared error
    pub mse: f64,
    /// Root mean squared error
    pub rmse: f64,
    /// Mean absolute error
    pub mae: f64,
}

/// Fits ridge regression for a single lambda value.
///
/// # Arguments
///
/// * `x` - Design matrix (n × p). Should include intercept column if `intercept=true`.
/// * `y` - Response vector (n elements)
/// * `options` - Ridge fitting options
///
/// # Returns
///
/// A [`RidgeFit`] containing the fit results.
///
/// # Errors
///
/// Returns an error if:
/// - `lambda < 0`
/// - Dimensions don't match
/// - Matrix is numerically singular
///
/// # Algorithm
///
/// Uses the augmented QR approach:
/// 1. Standardize X and center y (if requested)
/// 2. Build augmented system:
///    ```text
///    X_aug = [X_std; sqrt(lambda) * I_p]
///    y_aug = [y_centered; 0_p]
///    ```
/// 3. Solve using QR decomposition
/// 4. Unstandardize coefficients
///
/// # Example
///
/// ```rust,no_run
/// use linreg_core::linalg::Matrix;
/// use linreg_core::regularized::ridge::{ridge_fit, RidgeFitOptions};
///
/// let x = Matrix::new(3, 2, vec![
///     1.0, 2.0,
///     1.0, 3.0,
///     1.0, 4.0,
/// ]);
/// let y = vec![3.0, 5.0, 7.0];
///
/// let options = RidgeFitOptions {
///     lambda: 1.0,
///     intercept: true,
///     standardize: true,
/// };
///
/// let fit = ridge_fit(&x, &y, &options).unwrap();
/// println!("Intercept: {}", fit.intercept);
/// println!("Coefficients: {:?}", fit.coefficients);
/// ```
pub fn ridge_fit(x: &Matrix, y: &[f64], options: &RidgeFitOptions) -> Result<RidgeFit> {
    if options.lambda < 0.0 {
        return Err(Error::InvalidInput(
            "Lambda must be non-negative for ridge regression".to_string(),
        ));
    }

    let n = x.rows;
    let p = x.cols;

    if y.len() != n {
        return Err(Error::DimensionMismatch(
            format!("Length of y ({}) must match number of rows in X ({})", y.len(), n)
        ));
    }

    // Handle zero lambda: just do OLS
    if options.lambda == 0.0 {
        return ridge_ols_fit(x, y, options);
    }

    // Standardize X and center y
    let std_options = StandardizeOptions {
        intercept: options.intercept,
        standardize_x: options.standardize,
        standardize_y: false, // Don't standardize y for ridge
    };

    let (x_std, y_centered, std_info) = standardize_xy(x, y, &std_options);

    // Build augmented system: [X; sqrt(lambda)*I] * beta = [y; 0]
    // For the intercept column (if present), we don't add penalty
    let sqrt_lambda = options.lambda.sqrt();
    let intercept_col = if options.intercept { 1 } else { 0 };

    // Number of penalized coefficients (excluding intercept)
    let p_pen = p - intercept_col;

    // Augmented matrix dimensions
    let aug_n = n + p_pen;
    let aug_p = p;

    // Build augmented matrix
    let mut x_aug_data = vec![0.0; aug_n * aug_p];

    // Copy X_std to top portion
    for i in 0..n {
        for j in 0..p {
            x_aug_data[i * aug_p + j] = x_std.get(i, j);
        }
    }

    // Add sqrt(lambda) * I for penalized coefficients
    for i in 0..p_pen {
        let row = n + i;
        let col = intercept_col + i;
        x_aug_data[row * aug_p + col] = sqrt_lambda;
    }

    let x_aug = Matrix::new(aug_n, aug_p, x_aug_data);

    // Build augmented y vector
    let mut y_aug = vec![0.0; aug_n];
    for i in 0..n {
        y_aug[i] = y_centered[i];
    }
    // Remaining entries are already 0

    // Solve using QR decomposition
    let (q, r) = x_aug.qr();
    let beta_std = solve_upper_triangular_with_augmented_y(&r, &q, &y_aug, aug_n)?;

    // Unstandardize coefficients
    let (intercept, beta_orig) = unstandardize_coefficients(&beta_std, &std_info);

    // Compute fitted values and residuals on original scale
    let fitted = predict(x, intercept, &beta_orig);
    let residuals: Vec<f64> = y.iter().zip(fitted.iter()).map(|(yi, yh)| yi - yh).collect();

    // Compute effective degrees of freedom
    // For ridge: df = trace(X(X'X + lambda*I)^(-1)X')
    // This equals sum of eigenvalues / (eigenvalues + lambda)
    // We compute it using the hat matrix approach
    let df = compute_ridge_df(&x_std, options.lambda, intercept_col);

    // Compute model fit statistics
    let y_mean: f64 = y.iter().sum::<f64>() / n as f64;
    let ss_tot: f64 = y.iter().map(|yi| (yi - y_mean).powi(2)).sum();
    let ss_res: f64 = residuals.iter().map(|r| r.powi(2)).sum();
    let r_squared = if ss_tot > 1e-10 {
        1.0 - ss_res / ss_tot
    } else {
        1.0
    };

    // Adjusted R² using effective degrees of freedom
    let adj_r_squared = if ss_tot > 1e-10 && (n as f64) > df {
        1.0 - (1.0 - r_squared) * ((n - 1) as f64 / (n as f64 - df))
    } else {
        r_squared
    };

    let mse = ss_res / (n as f64 - 1.0); // Use n-1 for consistency
    let rmse = mse.sqrt();
    let mae: f64 = residuals.iter().map(|r| r.abs()).sum::<f64>() / n as f64;

    Ok(RidgeFit {
        lambda: options.lambda,
        intercept,
        coefficients: beta_orig,
        fitted_values: fitted,
        residuals,
        df,
        r_squared,
        adj_r_squared,
        mse,
        rmse,
        mae,
    })
}

/// Computes the effective degrees of freedom for ridge regression.
///
/// df = trace(H) where H = X(X'X + λI)^(-1)X'
///
/// For a QR decomposition of the standardized X, this equals the sum of
/// squared diagonal elements of R(R'R + λI)^(-1).
fn compute_ridge_df(x_std: &Matrix, lambda: f64, intercept_col: usize) -> f64 {
    let p = x_std.cols;

    // For small problems, compute directly
    if p <= 100 {
        // Get QR decomposition of X_std
        let (_q, _r) = x_std.qr();

        // Compute df = trace(X(X'X + λI)^(-1)X')
        // This equals trace(R(R'R + λI)^(-1)R') / n for centered data
        // A simpler approach: df = sum of (d_i^2 / (d_i^2 + lambda))
        // where d_i are singular values of X

        // For ridge, a simple approximation that works well:
        // df = sum_{j not penalized} 1 + sum_{j penalized} sigma_j^2 / (sigma_j^2 + lambda)
        // where sigma_j^2 are eigenvalues of X'X

        // Use the approximation: df ≈ p - lambda * trace((X'X + lambda*I)^(-1))
        // For now, use a simpler proxy
        let p_pen = p - intercept_col;
        let df_penalty = if lambda > 0.0 {
            // Approximate reduction in df due to penalty
            (p_pen as f64) * lambda / (1.0 + lambda)
        } else {
            0.0
        };

        (p as f64) - df_penalty
    } else {
        // For large p, use a simpler approximation
        p as f64 * lambda / (1.0 + lambda)
    }
}

/// OLS fit for lambda = 0 (special case of ridge).
fn ridge_ols_fit(x: &Matrix, y: &[f64], options: &RidgeFitOptions) -> Result<RidgeFit> {
    let n = x.rows;
    let p = x.cols;

    // Standardize (center) for consistency
    let std_options = StandardizeOptions {
        intercept: options.intercept,
        standardize_x: false,
        standardize_y: false,
    };

    let (_, _y_centered, std_info) = standardize_xy(x, y, &std_options);

    // Use QR decomposition for OLS
    let x_for_ols = if options.intercept {
        x.clone()
    } else {
        x.clone()
    };

    let (q, r) = x_for_ols.qr();
    let beta_std = solve_upper_triangular_with_augmented_y(&r, &q, y, n)?;

    // Unstandardize
    let (intercept, beta_orig) = unstandardize_coefficients(&beta_std, &std_info);

    // Compute fitted values and residuals
    let fitted = predict(x, intercept, &beta_orig);
    let residuals: Vec<f64> = y.iter().zip(fitted.iter()).map(|(yi, yh)| yi - yh).collect();

    // For OLS, df = p (or n - 1 if considering adjusted df)
    let df = p as f64;

    // Compute model fit statistics
    let y_mean: f64 = y.iter().sum::<f64>() / n as f64;
    let ss_tot: f64 = y.iter().map(|yi| (yi - y_mean).powi(2)).sum();
    let ss_res: f64 = residuals.iter().map(|r| r.powi(2)).sum();
    let r_squared = if ss_tot > 1e-10 {
        1.0 - ss_res / ss_tot
    } else {
        1.0
    };

    // Adjusted R² using effective degrees of freedom
    let adj_r_squared = if ss_tot > 1e-10 && n > p {
        1.0 - (1.0 - r_squared) * ((n - 1) as f64 / (n - p) as f64)
    } else {
        r_squared
    };

    let mse = ss_res / (n as f64 - p as f64);
    let rmse = mse.sqrt();
    let mae: f64 = residuals.iter().map(|r| r.abs()).sum::<f64>() / n as f64;

    Ok(RidgeFit {
        lambda: 0.0,
        intercept,
        coefficients: beta_orig,
        fitted_values: fitted,
        residuals,
        df,
        r_squared,
        adj_r_squared,
        mse,
        rmse,
        mae,
    })
}

/// Solves R * beta = Q^T * y_aug for beta.
///
/// This is a helper for the augmented QR approach.
fn solve_upper_triangular_with_augmented_y(
    r: &Matrix,
    q: &Matrix,
    y_aug: &[f64],
    aug_n: usize,
) -> Result<Vec<f64>> {
    let p = r.cols;

    // Compute Q^T * y_aug (only need first p rows since R is p × p or m × p)
    // Actually, Q is aug_n × aug_n, but we only need Q^T * y_aug for first p rows
    // since R has zeros below row p

    let mut qty = vec![0.0; p];

    // Compute Q^T * y_aug for the first p rows
    for i in 0..p {
        let mut sum = 0.0;
        for k in 0..aug_n {
            sum += q.get(k, i) * y_aug[k];
        }
        qty[i] = sum;
    }

    // Back substitution: solve R * beta = qty
    let mut beta = vec![0.0; p];

    for i in (0..p).rev() {
        let mut sum = qty[i];
        for j in (i + 1)..p {
            sum -= r.get(i, j) * beta[j];
        }

        let diag = r.get(i, i);
        if diag.abs() < 1e-14 {
            return Err(Error::ComputationFailed(
                "Matrix is singular to working precision".to_string(),
            ));
        }

        beta[i] = sum / diag;
    }

    Ok(beta)
}

/// Makes predictions using a ridge regression fit.
///
/// # Arguments
///
/// * `fit` - The ridge regression fit result
/// * `x_new` - New data matrix (n_new × p)
///
/// # Returns
///
/// Predictions for each row in x_new.
pub fn predict_ridge(fit: &RidgeFit, x_new: &Matrix) -> Vec<f64> {
    predict(x_new, fit.intercept, &fit.coefficients)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ridge_fit_simple() {
        // Simple test: perfect linear relationship
        let x_data = vec![
            1.0, 1.0,
            1.0, 2.0,
            1.0, 3.0,
            1.0, 4.0,
        ];
        let x = Matrix::new(4, 2, x_data);
        let y = vec![2.0, 4.0, 6.0, 8.0]; // y = 2 * x (with intercept 0)

        let options = RidgeFitOptions {
            lambda: 0.1,
            intercept: true,
            standardize: false,
        };

        let fit = ridge_fit(&x, &y, &options).unwrap();

        // With small lambda, should be close to OLS solution
        // OLS solution: intercept ≈ 0, slope ≈ 2
        assert!((fit.coefficients[1] - 2.0).abs() < 0.1);
        assert!(fit.intercept.abs() < 0.1);
    }

    #[test]
    fn test_ridge_fit_with_standardization() {
        let x_data = vec![
            1.0, 100.0,
            1.0, 200.0,
            1.0, 300.0,
            1.0, 400.0,
        ];
        let x = Matrix::new(4, 2, x_data);
        let y = vec![2.0, 4.0, 6.0, 8.0];

        let options = RidgeFitOptions {
            lambda: 1.0,
            intercept: true,
            standardize: true,
        };

        let fit = ridge_fit(&x, &y, &options).unwrap();

        // Predictions should be reasonable
        for i in 0..4 {
            assert!((fit.fitted_values[i] - y[i]).abs() < 2.0);
        }
    }

    #[test]
    fn test_ridge_zero_lambda_is_ols() {
        let x_data = vec![
            1.0, 1.0,
            1.0, 2.0,
            1.0, 3.0,
        ];
        let x = Matrix::new(3, 2, x_data);
        let y = vec![2.0, 4.0, 6.0];

        let options = RidgeFitOptions {
            lambda: 0.0,
            intercept: true,
            standardize: false,
        };

        let fit = ridge_fit(&x, &y, &options).unwrap();

        // Should be close to perfect fit for this data
        assert!((fit.fitted_values[0] - 2.0).abs() < 1e-6);
        assert!((fit.fitted_values[1] - 4.0).abs() < 1e-6);
        assert!((fit.fitted_values[2] - 6.0).abs() < 1e-6);
    }

    #[test]
    fn test_ridge_negative_lambda_error() {
        let x_data = vec![1.0, 1.0, 1.0, 2.0, 1.0, 3.0];
        let x = Matrix::new(3, 2, x_data);
        let y = vec![2.0, 4.0, 6.0];

        let options = RidgeFitOptions {
            lambda: -1.0,
            ..Default::default()
        };

        let result = ridge_fit(&x, &y, &options);
        assert!(result.is_err());
    }

    #[test]
    fn test_predict_ridge() {
        let x_data = vec![
            1.0, 1.0,
            1.0, 2.0,
            1.0, 3.0,
        ];
        let x = Matrix::new(3, 2, x_data);
        let y = vec![2.0, 4.0, 6.0];

        let options = RidgeFitOptions {
            lambda: 0.1,
            intercept: true,
            standardize: false,
        };

        let fit = ridge_fit(&x, &y, &options).unwrap();
        let preds = predict_ridge(&fit, &x);

        // Predictions on training data should equal fitted values
        for i in 0..3 {
            assert!((preds[i] - fit.fitted_values[i]).abs() < 1e-10);
        }
    }
}