strapdown/

linalg.rs

//! Linear algebra helpers for robust covariance square roots.
//!
//! Public API:
//!     pub fn matrix_square_root(matrix: &DMatrix<f64>) -> DMatrix<f64>
//!
//! Internal pipeline (each step isolated for testing):
//!     - symmetrize()
//!     - chol_sqrt()
//!     - chol_sqrt_with_jitter()
//!     - evd_symmetric_sqrt_with_floor()
//!
//! Strategy:
//! 1) Symmetrize P ← 0.5 (P + Pᵀ)
//! 2) Cholesky
//! 3) Jittered Cholesky (geometric ramp)
//! 4) Symmetric EVD with eigenvalue floor → S = U * sqrt(Λ⁺) * Uᵀ

use nalgebra::DMatrix;
use nalgebra::linalg::{Cholesky, SymmetricEigen};

/// Compute a robust symmetric square root `S` such that approximately `matrix ≈ S * Sᵀ`.
///
/// Attempts Cholesky decomposition first (yielding L such that matrix = L * L^T).
/// If Cholesky fails (e.g., matrix is not positive definite), it attempts to compute
/// the square root using eigenvalue decomposition (S = V * sqrt(D) * V^T).
///
/// # Arguments
/// * `matrix` - The DMatrix<f64> to find the square root of. It's assumed to be symmetric and square.
///
/// # Returns
/// * `Some(DMatrix<f64>)` containing a matrix square root.
///   The result from Cholesky is lower triangular. The result from eigenvalue decomposition is symmetric.
///   In both cases, if the result is `M`, then `matrix` approx `M * M.transpose()`.
/// * `None` if the matrix is not square or another fundamental issue prevents computation (though
///   this implementation tries to be robust for positive semi-definite cases).
pub fn matrix_square_root(matrix: &DMatrix<f64>) -> DMatrix<f64> {
    assert!(
        matrix.is_square(),
        "matrix_square_root: matrix must be square"
    );

    // Tunable guards (conservative defaults for double precision INS scales)
    const INITIAL_JITTER: f64 = 1e-12;
    const MAX_JITTER: f64 = 1e-6;
    const MAX_TRIES: usize = 6;
    const EIGEN_FLOOR: f64 = 1e-12;

    // 1) Symmetrize to kill round-off asymmetry
    let p = symmetrize(matrix);

    // 2) Cholesky (fast path)
    if let Some(s) = chol_sqrt(&p) {
        return s;
    }

    // 3) Jittered Cholesky
    if let Some(s) = chol_sqrt_with_jitter(&p, INITIAL_JITTER, MAX_JITTER, MAX_TRIES) {
        return s;
    }

    // 4) EVD fallback with eigenvalue floor — symmetric square root
    return evd_symmetric_sqrt_with_floor(&p, EIGEN_FLOOR);
}
/// Symmetrize a matrix: P ← 0.5 (P + Pᵀ)
///
/// Simple matrix symmetrization function that reduces round-off errors associated
/// with floating point arithmetic.
///
/// # Arguments
/// * `m` - the matrix to symmetrize
///
/// # Returns
/// A symmetrized version of the input matrix.
#[inline]
pub fn symmetrize(m: &DMatrix<f64>) -> DMatrix<f64> {
    0.5 * (m + m.transpose())
}
/// Plain Cholesky square root
///
/// Cholesky factorization that returns L such that P ≈ L Lᵀ, or None if it fails.
/// This is a quick way to initially attempt to calculate a matrix square root.
///
/// # Arguments
/// * ``p` - the matrix to factor
///
/// # Returns
/// A lower triangular matrix L such that P ≈ L Lᵀ, or None if it fails.
fn chol_sqrt(p: &DMatrix<f64>) -> Option<DMatrix<f64>> {
    Cholesky::new(p.clone()).map(|ch| ch.l().into_owned())
}
/// Cholesky with diagonal jitter (geometric ramp). Returns None if all tries fail.
///
/// Perform Cholesky decomposition with a jittered diagonal on a geometric ramp up.
/// Returns None if all tries fail.
fn chol_sqrt_with_jitter(
    p: &DMatrix<f64>,
    initial_jitter: f64,
    max_jitter: f64,
    max_tries: usize,
) -> Option<DMatrix<f64>> {
    let n = p.nrows();
    let mut jitter = initial_jitter;
    for _ in 0..max_tries {
        let mut pj = p.clone();
        for i in 0..n {
            pj[(i, i)] += jitter;
        }
        if let Some(ch) = Cholesky::new(pj) {
            return Some(ch.l().into_owned());
        }
        jitter *= 10.0;
        if jitter > max_jitter {
            break;
        }
    }
    None
}

/// Symmetric EVD square root with eigenvalue flooring:
/// S = U * sqrt(max(λ, floor)) * Uᵀ
fn evd_symmetric_sqrt_with_floor(p: &DMatrix<f64>, floor: f64) -> DMatrix<f64> {
    let se = SymmetricEigen::new(p.clone());
    let mut lambdas = se.eigenvalues;
    let u = se.eigenvectors;

    for i in 0..lambdas.len() {
        if lambdas[i] < floor {
            lambdas[i] = floor;
        }
    }

    let sqrt_vals = lambdas.map(|l| l.sqrt());
    let sigma_half = DMatrix::<f64>::from_diagonal(&sqrt_vals);
    &u * sigma_half * u.transpose()
}

#[derive(Debug, Clone, Copy)]
pub struct SolveOptions {
    pub initial_jitter: f64, // e.g., 1e-12
    pub max_jitter: f64,     // e.g., 1e-6
    pub max_tries: usize,    // e.g., 6
}

impl Default for SolveOptions {
    fn default() -> Self {
        Self {
            initial_jitter: 1e-12,
            max_jitter: 1e-6,
            max_tries: 6,
        }
    }
}
/// Solve A X = B for SPD-ish A via Cholesky, with jitter retries.
/// Returns None if all attempts fail.
pub fn chol_solve_spd(
    a: &DMatrix<f64>,
    b: &DMatrix<f64>,
    opt: SolveOptions,
) -> Option<DMatrix<f64>> {
    assert!(a.is_square(), "chol_solve_spd: A must be square");
    assert_eq!(a.nrows(), b.nrows(), "chol_solve_spd: A and B incompatible");

    // Symmetrize first (SPD drift is common).
    let a_sym = symmetrize(a);

    // Try plain Cholesky
    if let Some(ch) = Cholesky::new(a_sym.clone()) {
        return Some(ch.solve(b));
    }

    // Jitter ramp
    let n = a_sym.nrows();
    let mut jitter = opt.initial_jitter;
    for _ in 0..opt.max_tries {
        let mut a_j = a_sym.clone();
        for i in 0..n {
            a_j[(i, i)] += jitter;
        }
        if let Some(ch) = Cholesky::new(a_j) {
            return Some(ch.solve(b));
        }
        jitter *= 10.0;
        if jitter > opt.max_jitter {
            break;
        }
    }
    None
}

/// Robust SPD solve with sane defaults:
/// - Cholesky + jitter (preferred)
/// - Last resort: explicit inverse
pub fn robust_spd_solve(a: &DMatrix<f64>, b: &DMatrix<f64>) -> DMatrix<f64> {
    if let Some(x) = chol_solve_spd(a, b, SolveOptions::default()) {
        x
    } else if let Some(inv) = symmetrize(a).try_inverse() {
        &inv * b
    } else {
        panic!("robust_spd_solve: A is not invertible (even after jitter).");
    }
}

/* =============================== Tests ==================================== */

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

    fn approx_eq(a: &DMatrix<f64>, b: &DMatrix<f64>, tol: f64) -> bool {
        if a.shape() != b.shape() {
            return false;
        }
        let mut max_abs = 0.0f64;
        for i in 0..a.nrows() {
            for j in 0..a.ncols() {
                max_abs = max_abs.max((a[(i, j)] - b[(i, j)]).abs());
            }
        }
        max_abs <= tol
    }

    #[test]
    fn t_symmetrize() {
        let m = DMatrix::from_row_slice(2, 2, &[1.0, 2.0, 0.0, 3.0]);
        let s = symmetrize(&m);
        let s_expected = DMatrix::from_row_slice(2, 2, &[1.0, 1.0, 1.0, 3.0]);
        assert!(approx_eq(&s, &s_expected, 1e-15));
    }

    #[test]
    fn t_chol_sqrt_spd() {
        // P = A Aᵀ is SPD
        let a = DMatrix::from_row_slice(3, 3, &[1.0, 2.0, 0.5, 0.0, 1.0, -1.0, 0.0, 0.0, 0.2]);
        let p = &a * a.transpose();
        let s = chol_sqrt(&p).expect("Cholesky should succeed for SPD");
        let back = &s * s.transpose();
        assert!(approx_eq(&back, &p, 1e-12));
    }

    #[test]
    fn t_chol_sqrt_with_jitter() {
        // Nudge diagonal a hair negative to break plain Cholesky
        let a = DMatrix::from_row_slice(3, 3, &[1.0, 0.2, 0.0, 0.0, 1.0, 0.2, 0.0, 0.0, 1.0]);
        let mut p = &a * a.transpose();
        p[(2, 2)] -= 1e-10;

        //assert!(chol_sqrt(&p).is_none(), "plain Cholesky should fail here");
        let s =
            chol_sqrt_with_jitter(&p, 1e-12, 1e-6, 6).expect("jittered Cholesky should succeed");
        let back = &s * s.transpose();
        let p_sym = symmetrize(&p);
        assert!(approx_eq(&back, &p_sym, 1e-8));
    }

    #[test]
    fn t_evd_floor() {
        // Make P symmetric but with a negative eigenvalue, EVD should floor it.
        let p = DMatrix::from_row_slice(2, 2, &[0.0, 1.0, 1.0, 0.0]); // eigenvalues {+1, -1}
        let s = evd_symmetric_sqrt_with_floor(&p, 1e-12);
        let back = &s * s.transpose();
        // back should be PSD and close to symmetrized p with floor effects
        let p_sym = symmetrize(&p);
        assert_eq!(back.nrows(), p_sym.nrows());
        assert_eq!(back.ncols(), p_sym.ncols());
        // sanity: back is symmetric
        assert!(approx_eq(&back, &back.transpose(), 1e-14));
    }

    #[test]
    fn t_public_identity() {
        let i = DMatrix::<f64>::identity(4, 4);
        let s = matrix_square_root(&i);
        assert!(approx_eq(&s, &i, 1e-14));
        let back = &s * s.transpose();
        assert!(approx_eq(&back, &i, 1e-12));
    }

    #[test]
    fn t_public_nearly_spd() {
        let a = DMatrix::from_row_slice(3, 3, &[1.0, 0.1, 0.0, 0.0, 1.0, 0.2, 0.0, 0.0, 1.0]);
        let mut p = &a * a.transpose();
        p[(2, 2)] -= 1e-10;
        p[(0, 2)] += 1e-12; // asymmetry

        let s = matrix_square_root(&p);
        let back = &s * s.transpose();
        let p_sym = symmetrize(&p);
        assert!(approx_eq(&back, &p_sym, 1e-8));
    }

    #[test]
    #[should_panic]
    fn t_public_non_square_panics() {
        let m = DMatrix::<f64>::zeros(3, 2);
        let _ = matrix_square_root(&m);
    }
}

// ============ OLD ====================================

// Calculates a square root of a symmetric matrix.
//
// Attempts Cholesky decomposition first (yielding L such that matrix = L * L^T).
// If Cholesky fails (e.g., matrix is not positive definite), it attempts to compute
// the square root using eigenvalue decomposition (S = V * sqrt(D) * V^T).
// For eigenvalue decomposition, eigenvalues are clamped to be non-negative.
//
// # Arguments
// * `matrix` - The DMatrix<f64> to find the square root of. It's assumed to be symmetric and square.
//
// # Returns
// * `Some(DMatrix<f64>)` containing a matrix square root.
//   The result from Cholesky is lower triangular. The result from eigenvalue decomposition is symmetric.
//   In both cases, if the result is `M`, then `matrix` approx `M * M.transpose()`.
// * `None` if the matrix is not square or another fundamental issue prevents computation (though
//   this implementation tries to be robust for positive semi-definite cases).
//pub fn matrix_square_root(matrix: &DMatrix<f64>) -> DMatrix<f64> {
//    if !matrix.is_square() {
//        panic!("Error: Matrix must be square to compute square root.");
//    }
//    // Attempt Cholesky decomposition (yields L where matrix = L * L^T)
//    // Cholesky requires the matrix to be symmetric positive definite.
//    match cholesky_pass(matrix) {
//        Some(chol_l) => {
//            return chol_l;
//        }
//        None => {
//            //println!("Cholesky decomposition failed. Attempting eigenvalue decomposition.");
//        }
//    }
//    // If Cholesky failed, we try eigenvalue decomposition.
//    match eigenvalue_pass(matrix) {
//        Some(eigen_sqrt) => eigen_sqrt,
//        None => {
//            panic!(
//                "Cholesky and Eigenvalue decomposition failed. No valid square root found for the covariance matrix: \n {:?}",
//                matrix
//            );
//        }
//    }
//}
// Attempts to compute the matrix square root using Cholesky decomposition.
//
// This method is only applicable to symmetric positive definite matrices.
// If successful, it returns the lower triangular matrix `L` such that `matrix = L * L.transpose()`.
//
// When the computation _fails_ (e.g., the matrix is not positive definite or not square),
// a None value is returned instead of panicking, permitting the public API to proceed to the
// next method.
//
// # Arguments
// * `matrix` - The DMatrix<f64> to find the square root of. Assumed to be symmetric and square.
//
// # Returns
// * `Some(DMatrix<f64>)` containing the lower triangular Cholesky factor `L`.
// * `None` if the matrix is not positive definite or not square.
// fn cholesky_pass(matrix: &DMatrix<f64>) -> Option<DMatrix<f64>> {
//     if !matrix.is_square() {
//         eprintln!("Error: Matrix must be square for Cholesky decomposition.");
//         return None;
//     }
//     matrix
//         .clone()
//         .cholesky()
//         .map(|chol: Cholesky<f64, nalgebra::Dyn>| chol.l())
// }
// Computes a symmetric matrix square root using eigenvalue decomposition.
//
// This method is suitable for symmetric positive semi-definite matrices.
// It returns a symmetric matrix `S` such that `matrix = S * S`.
// Eigenvalues are clamped to be non-negative to handle positive semi-definite cases
// and minor numerical inaccuracies.
//
// When the computation _fails_ (e.g., the matrix is not positive definite or not square),
// a None value is returned instead of panicking, permitting the public API to proceed to the
// next method.
//
// # Arguments
// * `matrix` - The DMatrix<f64> to find the square root of. Assumed to be symmetric and square.
//
// # Returns
// * `Some(DMatrix<f64>)` containing the symmetric matrix square root `S`.
// * `None` if the matrix is not square (though this should be checked by the caller for symmetry assumptions).
// fn eigenvalue_pass(matrix: &DMatrix<f64>) -> Option<DMatrix<f64>> {
//     if !matrix.is_square() {
//         eprintln!("Error: Matrix must be square for eigenvalue decomposition based square root.");
//         return None;
//     }
//     // For eigenvalue decomposition of a symmetric matrix,
//     // we use `symmetric_eigen`. This returns real eigenvalues and orthogonal eigenvectors.
//     let eigen_decomposition: SymmetricEigen<f64, nalgebra::Dyn> = matrix.clone().symmetric_eigen();
//     let eigenvalues = eigen_decomposition.eigenvalues;
//     let eigenvectors = eigen_decomposition.eigenvectors;
//
//     // Check for significantly negative eigenvalues, indicating non-positive semi-definiteness.
//     // While we clamp them, a warning is useful for diagnosis.
//     if eigenvalues.iter().any(|&val| val < -1e-9) {
//         println!(
//             "Warning: Negative eigenvalues encountered during eigenvalue decomposition. The input matrix was not positive semi-definite."
//         );
//     //     println!("{:?}", matrix.data);
//     //     // return None;
//     }
//
//     // Create diagonal matrix of sqrt(eigenvalues), clamping eigenvalues to be non-negative.
//     // `DMatrix::from_diagonal` takes a DVector.
//     let sqrt_eigenvalues_diag_vec = eigenvalues.map(|val| val.max(1e-9).sqrt());
//     let sqrt_eigenvalues_diag = DMatrix::from_diagonal(&sqrt_eigenvalues_diag_vec);
//
//     // Reconstruct the square root: S = V * sqrt(D) * V^T
//     // This S will be symmetric, and S * S = matrix (or S * S^T = matrix).
//     let sqrt_m = eigenvectors.clone() * sqrt_eigenvalues_diag * eigenvectors.transpose();
//
//     Some(sqrt_m)
// }
//
// #[cfg(test)]
// mod tests {
//     use super::*;
//     use nalgebra::DMatrix;
//     use std::sync::LazyLock;
//
//     static BASIC_SQRT: LazyLock<DMatrix<f64>> = LazyLock::new(|| {
//         DMatrix::from_row_slice(3, 3, &[4.0, 0.0, 0.0, 0.0, 9.0, 0.0, 0.0, 0.0, 16.0])
//     });
//     static POSITIVE_DEFINITE: LazyLock<DMatrix<f64>> = LazyLock::new(|| {
//         DMatrix::from_row_slice(3, 3, &[4.0, 2.0, 0.0, 2.0, 9.0, 3.0, 0.0, 3.0, 16.0])
//     });
//     static POSITIVE_SEMI_DEFINITE: LazyLock<DMatrix<f64>> = LazyLock::new(|| {
//         DMatrix::from_row_slice(3, 3, &[1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0])
//     });
//     static NEGATIVE_DEFINITE: LazyLock<DMatrix<f64>> = LazyLock::new(|| {
//         DMatrix::from_row_slice(3, 3, &[-4.0, 0.0, 0.0, 0.0, -9.0, 0.0, 0.0, 0.0, -16.0])
//     });
//     static NEGATIVE_SEMI_DEFINITE: LazyLock<DMatrix<f64>> = LazyLock::new(|| {
//         DMatrix::from_row_slice(3, 3, &[-1.0, 0.0, -1.0, 0.0, -1.0, 0.0, -1.0, 0.0, -1.0])
//     });
//     static NON_SQUARE: LazyLock<DMatrix<f64>> =
//         LazyLock::new(|| DMatrix::from_row_slice(2, 3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]));
//
//     /// Helper function to verify if a matrix is a valid square root of another matrix.
//     /// Returns true if sqrt_matrix * sqrt_matrix.T ≈ original_matrix within tolerance.
//     fn is_valid_square_root(
//         sqrt_matrix: &DMatrix<f64>,
//         original_matrix: &DMatrix<f64>,
//         tolerance: f64,
//     ) -> bool {
//         let reconstructed = sqrt_matrix * sqrt_matrix.transpose();
//
//         if reconstructed.nrows() != original_matrix.nrows()
//             || reconstructed.ncols() != original_matrix.ncols()
//         {
//             return false;
//         }
//
//         for i in 0..original_matrix.nrows() {
//             for j in 0..original_matrix.ncols() {
//                 if (reconstructed[(i, j)] - original_matrix[(i, j)]).abs() > tolerance {
//                     return false;
//                 }
//             }
//         }
//         true
//     }
//     // Test matrix square root calculation
//     #[test]
//     fn cholesky_square_root() {
//         let sqrt_matrix = matrix_square_root(&BASIC_SQRT);
//         assert!(is_valid_square_root(&sqrt_matrix, &BASIC_SQRT, 1e-9));
//     }
//     #[test]
//     fn cholesky_positive_definite() {
//         let sqrt_matrix = matrix_square_root(&POSITIVE_DEFINITE);
//         assert!(is_valid_square_root(&sqrt_matrix, &POSITIVE_DEFINITE, 1e-9));
//     }
//     #[test]
//     #[should_panic]
//     fn cholesky_negative_definite() {
//         // This should panic because the matrix is negative definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn cholesky_negative_semi_definite() {
//         // This should panic because the matrix is negative semi-definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_SEMI_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn cholesky_non_square() {
//         // This should panic because the matrix is not square.
//         let _sqrt_matrix = matrix_square_root(&NON_SQUARE);
//     }
//     #[test]
//     fn eigenvalue_square_root() {
//         let sqrt_matrix = matrix_square_root(&POSITIVE_SEMI_DEFINITE);
//         assert!(is_valid_square_root(
//             &sqrt_matrix,
//             &POSITIVE_SEMI_DEFINITE,
//             1e-9
//         ));
//     }
//     #[test]
//     fn eigenvalue_positive_definite() {
//         let sqrt_matrix = matrix_square_root(&POSITIVE_DEFINITE);
//         assert!(is_valid_square_root(&sqrt_matrix, &POSITIVE_DEFINITE, 1e-9));
//     }
//     #[test]
//     fn eigenvalue_positive_semi_definite() {
//         let sqrt_matrix = matrix_square_root(&POSITIVE_SEMI_DEFINITE);
//         assert!(is_valid_square_root(
//             &sqrt_matrix,
//             &POSITIVE_SEMI_DEFINITE,
//             1e-9
//         ));
//     }
//     #[test]
//     #[should_panic]
//     fn eigenvalue_negative_definite() {
//         // This should panic because the matrix is negative definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn eigenvalue_negative_semi_definite() {
//         // This should panic because the matrix is negative semi-definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_SEMI_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn eigenvalue_non_square() {
//         // This should panic because the matrix is not square.
//         let _sqrt_matrix = matrix_square_root(&NON_SQUARE);
//     }
//     #[test]
//     fn public_api_square_root() {
//         let sqrt_matrix = matrix_square_root(&POSITIVE_DEFINITE);
//         assert!(is_valid_square_root(&sqrt_matrix, &POSITIVE_DEFINITE, 1e-9));
//         let sqrt_matrix = matrix_square_root(&POSITIVE_SEMI_DEFINITE);
//         assert!(is_valid_square_root(
//             &sqrt_matrix,
//             &POSITIVE_SEMI_DEFINITE,
//             1e-9
//         ));
//         let sqrt_matrix = matrix_square_root(&BASIC_SQRT);
//         assert!(is_valid_square_root(&sqrt_matrix, &BASIC_SQRT, 1e-9));
//     }
//     #[test]
//     #[should_panic]
//     fn public_api_negative_definite() {
//         // This should panic because the matrix is negative definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn public_api_negative_semi_definite() {
//         // This should panic because the matrix is negative semi-definite.
//         let _sqrt_matrix = matrix_square_root(&NEGATIVE_SEMI_DEFINITE);
//     }
//     #[test]
//     #[should_panic]
//     fn public_api_non_square() {
//         // This should panic because the matrix is not square.
//         let _sqrt_matrix = matrix_square_root(&NON_SQUARE);
//     }
// }
//