berkus
/
wifi-densepose
mirror of https://github.com/ruvnet/wifi-densepose.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
wifi-densepose/vendor/ruvector/crates/ruvector-solver/src/cg.rs

1100 lines
36 KiB
Rust
Raw Blame History

//! Conjugate Gradient solver for symmetric positive-definite systems.
//!
//! Solves `Ax = b` where `A` is a symmetric positive-definite (SPD) sparse
//! matrix in CSR format. The algorithm converges in at most `n` iterations
//! for an `n x n` system, but in practice converges in
//! `O(sqrt(kappa) * log(1/eps))` iterations where `kappa = cond(A)`.
//!
//! # Algorithm
//!
//! Implements the Hestenes-Stiefel variant of Conjugate Gradient:
//!
//! ```text
//! r = b - A*x
//! z = M^{-1} * r (preconditioner; z = r when disabled)
//! p = z
//! rz = r . z
//!
//! for k in 0..max_iterations:
//! Ap = A * p
//! alpha = rz / (p . Ap)
//! x = x + alpha * p
//! r = r - alpha * Ap
//! if ||r||_2 < tolerance * ||b||_2:
//! converged; break
//! z = M^{-1} * r
//! rz_new = r . z
//! beta = rz_new / rz
//! p = z + beta * p
//! rz = rz_new
//! ```
//!
//! # Preconditioning
//!
//! When `use_preconditioner` is `true`, a diagonal (Jacobi) preconditioner is
//! applied: `M = diag(A)`, so that `z_i = r_i / A_{ii}`. This reduces the
//! effective condition number for diagonally-dominant systems without adding
//! significant per-iteration cost.
//!
//! # Numerical precision
//!
//! All dot products and norm computations use `f64` accumulation even though
//! the matrix may store `f32` values, preventing catastrophic cancellation in
//! the inner products that drive the CG recurrence.
//!
//! # Convergence
//!
//! Theoretical bound:
//! `||x_k - x*||_A <= 2 * ((sqrt(kappa) - 1)/(sqrt(kappa) + 1))^k * ||x_0 - x*||_A`
//! where `kappa` is the 2-condition number of `A`.
use std::time::Instant;
use tracing::{debug, trace, warn};
use crate::error::{SolverError, ValidationError};
use crate::traits::SolverEngine;
use crate::types::{
Algorithm, ComplexityClass, ComplexityEstimate, ComputeBudget, ConvergenceInfo, CsrMatrix,
SolverResult, SparsityProfile,
};
// ═══════════════════════════════════════════════════════════════════════════
// Helper functions -- f64-accumulated linear algebra primitives
// ═══════════════════════════════════════════════════════════════════════════
/// Compute the dot product of two `f32` slices using `f64` accumulation.
///
/// Uses a 4-wide accumulator to exploit instruction-level parallelism and
/// reduce the dependency chain length, preventing precision loss in the
/// inner products that drive the CG recurrence.
///
/// # Panics
///
/// Debug-asserts that `a.len() == b.len()`.
#[inline]
pub fn dot_product_f64(a: &[f32], b: &[f32]) -> f64 {
assert_eq!(a.len(), b.len(), "dot_product_f64: length mismatch");
let n = a.len();
let chunks = n / 4;
let remainder = n % 4;
let mut acc0: f64 = 0.0;
let mut acc1: f64 = 0.0;
let mut acc2: f64 = 0.0;
let mut acc3: f64 = 0.0;
for i in 0..chunks {
let j = i * 4;
acc0 += a[j] as f64 * b[j] as f64;
acc1 += a[j + 1] as f64 * b[j + 1] as f64;
acc2 += a[j + 2] as f64 * b[j + 2] as f64;
acc3 += a[j + 3] as f64 * b[j + 3] as f64;
}
let base = chunks * 4;
for i in 0..remainder {
acc0 += a[base + i] as f64 * b[base + i] as f64;
}
(acc0 + acc1) + (acc2 + acc3)
}
/// Compute the dot product of two `f64` slices with 4-wide accumulation.
///
/// Used internally when the working vectors are already `f64`.
#[inline]
fn dot_f64(a: &[f64], b: &[f64]) -> f64 {
assert_eq!(a.len(), b.len(), "dot_f64: length mismatch");
let n = a.len();
let chunks = n / 4;
let remainder = n % 4;
let mut acc0: f64 = 0.0;
let mut acc1: f64 = 0.0;
let mut acc2: f64 = 0.0;
let mut acc3: f64 = 0.0;
for i in 0..chunks {
let j = i * 4;
acc0 += a[j] * b[j];
acc1 += a[j + 1] * b[j + 1];
acc2 += a[j + 2] * b[j + 2];
acc3 += a[j + 3] * b[j + 3];
}
let base = chunks * 4;
for i in 0..remainder {
acc0 += a[base + i] * b[base + i];
}
(acc0 + acc1) + (acc2 + acc3)
}
/// Compute `y[i] += alpha * x[i]` for all `i` (AXPY operation, `f32`).
///
/// # Panics
///
/// Debug-asserts that `x.len() == y.len()`.
#[inline]
pub fn axpy(alpha: f32, x: &[f32], y: &mut [f32]) {
assert_eq!(x.len(), y.len(), "axpy: length mismatch");
let n = x.len();
let chunks = n / 4;
let base = chunks * 4;
for i in 0..chunks {
let j = i * 4;
y[j] += alpha * x[j];
y[j + 1] += alpha * x[j + 1];
y[j + 2] += alpha * x[j + 2];
y[j + 3] += alpha * x[j + 3];
}
for i in base..n {
y[i] += alpha * x[i];
}
}
/// Compute `y[i] += alpha * x[i]` for all `i` (AXPY operation, `f64`).
#[inline]
fn axpy_f64(alpha: f64, x: &[f64], y: &mut [f64]) {
assert_eq!(x.len(), y.len(), "axpy_f64: length mismatch");
let n = x.len();
let chunks = n / 4;
let base = chunks * 4;
for i in 0..chunks {
let j = i * 4;
y[j] += alpha * x[j];
y[j + 1] += alpha * x[j + 1];
y[j + 2] += alpha * x[j + 2];
y[j + 3] += alpha * x[j + 3];
}
for i in base..n {
y[i] += alpha * x[i];
}
}
/// Compute the L2 norm of an `f32` slice using `f64` accumulation.
///
/// Returns `sqrt(sum(x_i^2))` computed entirely in `f64`.
#[inline]
pub fn norm2(x: &[f32]) -> f64 {
dot_product_f64(x, x).sqrt()
}
/// Compute the L2 norm of an `f64` slice.
#[inline]
fn norm2_f64(x: &[f64]) -> f64 {
dot_f64(x, x).sqrt()
}
// ═══════════════════════════════════════════════════════════════════════════
// ConjugateGradientSolver
// ═══════════════════════════════════════════════════════════════════════════
/// Conjugate Gradient solver for symmetric positive-definite sparse systems.
///
/// Stores the solver configuration (tolerance, iteration cap, preconditioning).
/// The solve itself is stateless and may be invoked concurrently on different
/// inputs from multiple threads.
#[derive(Debug, Clone)]
pub struct ConjugateGradientSolver {
/// Relative residual convergence tolerance.
///
/// The solver declares convergence when `||r||_2 < tolerance * ||b||_2`.
tolerance: f64,
/// Maximum number of CG iterations before declaring non-convergence.
max_iterations: usize,
/// Whether to apply diagonal (Jacobi) preconditioning.
///
/// When `true`, the preconditioner `M = diag(A)` is used, computing
/// `z_i = r_i / A_{ii}` each iteration. Beneficial for diagonally-dominant
/// systems at the cost of O(n) extra work per iteration.
use_preconditioner: bool,
}
impl ConjugateGradientSolver {
/// Create a new CG solver.
///
/// # Arguments
///
/// * `tolerance` -- Relative residual threshold for convergence. Must be
/// positive and finite.
/// * `max_iterations` -- Upper bound on CG iterations. Must be >= 1.
/// * `use_preconditioner` -- Enable diagonal (Jacobi) preconditioning.
pub fn new(tolerance: f64, max_iterations: usize, use_preconditioner: bool) -> Self {
Self {
tolerance,
max_iterations,
use_preconditioner,
}
}
/// Return the configured tolerance.
#[inline]
pub fn tolerance(&self) -> f64 {
self.tolerance
}
/// Return the configured maximum iterations.
#[inline]
pub fn max_iterations(&self) -> usize {
self.max_iterations
}
/// Return whether preconditioning is enabled.
#[inline]
pub fn use_preconditioner(&self) -> bool {
self.use_preconditioner
}
// -------------------------------------------------------------------
// Input validation
// -------------------------------------------------------------------
/// Validate inputs before entering the CG loop.
fn validate(&self, matrix: &CsrMatrix<f64>, rhs: &[f64]) -> Result<(), SolverError> {
if matrix.rows != matrix.cols {
return Err(SolverError::InvalidInput(
ValidationError::DimensionMismatch(format!(
"CG requires a square matrix but got {}x{}",
matrix.rows, matrix.cols,
)),
));
}
if rhs.len() != matrix.rows {
return Err(SolverError::InvalidInput(
ValidationError::DimensionMismatch(format!(
"rhs length {} does not match matrix rows {}",
rhs.len(),
matrix.rows,
)),
));
}
if matrix.row_ptr.len() != matrix.rows + 1 {
return Err(SolverError::InvalidInput(
ValidationError::DimensionMismatch(format!(
"row_ptr length {} does not equal rows + 1 = {}",
matrix.row_ptr.len(),
matrix.rows + 1,
)),
));
}
if !self.tolerance.is_finite() || self.tolerance <= 0.0 {
return Err(SolverError::InvalidInput(
ValidationError::ParameterOutOfRange {
name: "tolerance".into(),
value: self.tolerance.to_string(),
expected: "positive finite value".into(),
},
));
}
if self.max_iterations == 0 {
return Err(SolverError::InvalidInput(
ValidationError::ParameterOutOfRange {
name: "max_iterations".into(),
value: "0".into(),
expected: ">= 1".into(),
},
));
}
Ok(())
}
// -------------------------------------------------------------------
// Jacobi preconditioner
// -------------------------------------------------------------------
/// Build the Jacobi (diagonal) preconditioner from `A`.
///
/// Returns `inv_diag[i] = 1.0 / A_{ii}`. Zero or near-zero diagonal
/// entries are replaced with `1.0` to prevent division by zero.
fn build_jacobi_preconditioner(matrix: &CsrMatrix<f64>) -> Vec<f64> {
let n = matrix.rows;
let mut inv_diag = vec![1.0f64; n];
for row in 0..n {
let start = matrix.row_ptr[row];
let end = matrix.row_ptr[row + 1];
for idx in start..end {
if matrix.col_indices[idx] == row {
let diag_val = matrix.values[idx];
if diag_val.abs() > f64::EPSILON {
inv_diag[row] = 1.0 / diag_val;
}
break;
}
}
}
inv_diag
}
/// Apply the diagonal preconditioner: `z[i] = inv_diag[i] * r[i]`.
#[inline]
fn apply_preconditioner(inv_diag: &[f64], r: &[f64], z: &mut [f64]) {
assert_eq!(inv_diag.len(), r.len());
assert_eq!(r.len(), z.len());
let n = r.len();
let chunks = n / 4;
let base = chunks * 4;
for i in 0..chunks {
let j = i * 4;
z[j] = inv_diag[j] * r[j];
z[j + 1] = inv_diag[j + 1] * r[j + 1];
z[j + 2] = inv_diag[j + 2] * r[j + 2];
z[j + 3] = inv_diag[j + 3] * r[j + 3];
}
for i in base..n {
z[i] = inv_diag[i] * r[i];
}
}
// -------------------------------------------------------------------
// Core CG algorithm
// -------------------------------------------------------------------
/// Core CG algorithm implementation.
///
/// Works entirely in `f64` precision internally. The final solution is
/// down-cast to `f32` in the returned [`SolverResult`] to match the
/// crate's output contract.
fn solve_inner(
&self,
matrix: &CsrMatrix<f64>,
rhs: &[f64],
budget: &ComputeBudget,
) -> Result<SolverResult, SolverError> {
let start_time = Instant::now();
let n = matrix.rows;
// Effective limits: take the tighter of our config and the budget.
let effective_max_iter = self.max_iterations.min(budget.max_iterations);
let effective_tol = self.tolerance.min(budget.tolerance);
// --- Trivial case: zero-dimensional system ---
if n == 0 {
return Ok(SolverResult {
solution: vec![],
iterations: 0,
residual_norm: 0.0,
wall_time: start_time.elapsed(),
convergence_history: vec![],
algorithm: Algorithm::CG,
});
}
// --- Allocate working vectors (all f64) ---
let mut x = vec![0.0f64; n]; // solution (initial guess = zero)
let mut r = vec![0.0f64; n]; // residual
let mut z = vec![0.0f64; n]; // preconditioned residual
let mut p = vec![0.0f64; n]; // search direction
let mut ap = vec![0.0f64; n]; // A * p scratch buffer
// --- Build preconditioner (if enabled) ---
let inv_diag = if self.use_preconditioner {
Some(Self::build_jacobi_preconditioner(matrix))
} else {
None
};
// --- r = b - A*x. Since x = 0, r = b ---
r.copy_from_slice(rhs);
// --- Convergence threshold: ||r||_2 < tol * ||b||_2 (relative) ---
let b_norm = norm2_f64(rhs);
let abs_tolerance = effective_tol * b_norm;
// Handle zero RHS: the solution is the zero vector.
if b_norm < f64::EPSILON {
debug!("CG: zero RHS detected, returning zero solution");
return Ok(SolverResult {
solution: vec![0.0f32; n],
iterations: 0,
residual_norm: 0.0,
wall_time: start_time.elapsed(),
convergence_history: vec![],
algorithm: Algorithm::CG,
});
}
let initial_residual_norm = norm2_f64(&r);
// --- z = M^{-1} * r ---
match &inv_diag {
Some(diag) => Self::apply_preconditioner(diag, &r, &mut z),
None => z.copy_from_slice(&r),
}
// --- p = z ---
p.copy_from_slice(&z);
// --- rz = r . z ---
let mut rz = dot_f64(&r, &z);
let mut convergence_history = Vec::with_capacity(effective_max_iter.min(256));
let mut converged = false;
debug!(
"CG: n={}, nnz={}, tol={:.2e}, max_iter={}, precond={}",
n,
matrix.nnz(),
effective_tol,
effective_max_iter,
self.use_preconditioner,
);
// ===============================================================
// Main CG loop (Hestenes-Stiefel)
// ===============================================================
for k in 0..effective_max_iter {
// --- Budget: wall-time check ---
if start_time.elapsed() > budget.max_time {
warn!("CG: wall-time budget exhausted at iteration {k}");
return Err(SolverError::BudgetExhausted {
reason: format!(
"wall-time limit {:?} exceeded at iteration {k}",
budget.max_time,
),
elapsed: start_time.elapsed(),
});
}
// --- Ap = A * p (sparse matrix-vector product) ---
matrix.spmv(&p, &mut ap);
// --- alpha = rz / (p . Ap) ---
let p_dot_ap = dot_f64(&p, &ap);
// Guard: if p.Ap <= 0 the matrix is not SPD or we hit numerical
// breakdown.
if p_dot_ap <= 0.0 {
warn!("CG: non-positive p.Ap = {p_dot_ap:.4e} at iteration {k}");
return Err(SolverError::NumericalInstability {
iteration: k,
detail: format!("p.Ap = {p_dot_ap:.6e} <= 0; matrix may not be SPD",),
});
}
let alpha = rz / p_dot_ap;
// --- x = x + alpha * p ---
axpy_f64(alpha, &p, &mut x);
// --- r = r - alpha * Ap ---
axpy_f64(-alpha, &ap, &mut r);
// --- Convergence check: ||r||_2 < tol * ||b||_2 ---
let r_norm = norm2_f64(&r);
convergence_history.push(ConvergenceInfo {
iteration: k,
residual_norm: r_norm,
});
trace!(
"CG iter {k}: ||r|| = {r_norm:.6e}, rel = {:.6e}",
r_norm / b_norm,
);
if r_norm < abs_tolerance {
converged = true;
debug!(
"CG converged at iteration {k}: ||r|| = {r_norm:.6e}, \
rel = {:.6e}",
r_norm / b_norm,
);
break;
}
// --- Divergence detection ---
// If ||r|| has grown by 10x from the initial residual, the
// system is likely indefinite or ill-conditioned beyond rescue.
if r_norm > 10.0 * initial_residual_norm {
warn!(
"CG: divergence at iteration {k}: ||r|| = {r_norm:.6e} \
> 10 * ||r_0|| = {:.6e}",
10.0 * initial_residual_norm,
);
return Err(SolverError::NumericalInstability {
iteration: k,
detail: format!(
"residual diverged: ||r|| = {r_norm:.6e} exceeds \
10x initial residual {initial_residual_norm:.6e}",
),
});
}
// --- z = M^{-1} * r ---
match &inv_diag {
Some(diag) => Self::apply_preconditioner(diag, &r, &mut z),
None => z.copy_from_slice(&r),
}
// --- rz_new = r . z ---
let rz_new = dot_f64(&r, &z);
// Guard: stagnation when rz is near-zero.
if rz.abs() < f64::EPSILON * f64::EPSILON {
warn!("CG: rz near zero at iteration {k}, stagnation");
return Err(SolverError::NumericalInstability {
iteration: k,
detail: format!("rz = {rz:.6e} is near zero; solver stagnated",),
});
}
// --- beta = rz_new / rz ---
let beta = rz_new / rz;
// --- p = z + beta * p ---
for i in 0..n {
p[i] = z[i] + beta * p[i];
}
// --- rz = rz_new ---
rz = rz_new;
}
let wall_time = start_time.elapsed();
let final_residual = norm2_f64(&r);
if !converged {
debug!(
"CG: non-convergence after {} iterations, ||r|| = {final_residual:.6e}",
effective_max_iter,
);
return Err(SolverError::NonConvergence {
iterations: effective_max_iter,
residual: final_residual,
tolerance: abs_tolerance,
});
}
// Down-cast the f64 solution to f32 for SolverResult.
let solution_f32: Vec<f32> = x.iter().map(|&v| v as f32).collect();
Ok(SolverResult {
solution: solution_f32,
iterations: convergence_history.len(),
residual_norm: final_residual,
wall_time,
convergence_history,
algorithm: Algorithm::CG,
})
}
}
// ═══════════════════════════════════════════════════════════════════════════
// SolverEngine trait implementation
// ═══════════════════════════════════════════════════════════════════════════
impl SolverEngine for ConjugateGradientSolver {
/// Solve `Ax = b` using the Conjugate Gradient method.
///
/// # Errors
///
/// * [`SolverError::InvalidInput`] -- dimension mismatch or invalid params.
/// * [`SolverError::NumericalInstability`] -- divergence or non-SPD matrix.
/// * [`SolverError::NonConvergence`] -- iteration limit exceeded.
/// * [`SolverError::BudgetExhausted`] -- wall-time limit exceeded.
fn solve(
&self,
matrix: &CsrMatrix<f64>,
rhs: &[f64],
budget: &ComputeBudget,
) -> Result<SolverResult, SolverError> {
self.validate(matrix, rhs)?;
self.solve_inner(matrix, rhs, budget)
}
/// Estimate CG complexity from the sparsity profile.
///
/// CG converges in `O(sqrt(kappa))` iterations, each costing `O(nnz)` for
/// the SpMV plus `O(n)` for the vector updates.
fn estimate_complexity(&self, profile: &SparsityProfile, n: usize) -> ComplexityEstimate {
// Estimated iterations from condition number, clamped to max_iterations.
let est_iters = (profile.estimated_condition.sqrt() as usize)
.max(1)
.min(self.max_iterations);
// FLOPs per iteration: 2*nnz (SpMV) + 6*n (dot products, axpy ops).
let flops_per_iter = 2 * profile.nnz as u64 + 6 * n as u64;
let estimated_flops = est_iters as u64 * flops_per_iter;
// Memory: 5 vectors of length n (x, r, z, p, Ap) plus preconditioner.
let vec_bytes = n * std::mem::size_of::<f64>();
let precond_bytes = if self.use_preconditioner {
vec_bytes
} else {
0
};
let estimated_memory_bytes = 5 * vec_bytes + precond_bytes;
ComplexityEstimate {
algorithm: Algorithm::CG,
estimated_flops,
estimated_iterations: est_iters,
estimated_memory_bytes,
complexity_class: ComplexityClass::SqrtCondition,
}
}
/// Return the algorithm identifier.
fn algorithm(&self) -> Algorithm {
Algorithm::CG
}
}
// ═══════════════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════════════
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
/// Build a symmetric tridiagonal SPD matrix (f64):
/// diag = 4.0, off-diag = -1.0
fn tridiagonal_spd(n: usize) -> CsrMatrix<f64> {
let mut entries = Vec::with_capacity(3 * n);
for i in 0..n {
if i > 0 {
entries.push((i, i - 1, -1.0f64));
}
entries.push((i, i, 4.0f64));
if i + 1 < n {
entries.push((i, i + 1, -1.0f64));
}
}
CsrMatrix::<f64>::from_coo(n, n, entries)
}
/// Build a diagonal matrix from the given values.
fn diagonal_matrix(diag: &[f64]) -> CsrMatrix<f64> {
let n = diag.len();
let entries: Vec<_> = diag.iter().enumerate().map(|(i, &v)| (i, i, v)).collect();
CsrMatrix::<f64>::from_coo(n, n, entries)
}
/// Build an identity matrix of dimension n.
fn identity(n: usize) -> CsrMatrix<f64> {
CsrMatrix::<f64>::identity(n)
}
fn default_budget() -> ComputeBudget {
ComputeBudget {
max_time: Duration::from_secs(30),
max_iterations: 10_000,
tolerance: 1e-10,
}
}
// -----------------------------------------------------------------
// dot_product_f64
// -----------------------------------------------------------------
#[test]
fn dot_product_f64_basic() {
let a = vec![1.0f32, 2.0, 3.0];
let b = vec![4.0f32, 5.0, 6.0];
let result = dot_product_f64(&a, &b);
assert!((result - 32.0).abs() < 1e-10);
}
#[test]
fn dot_product_f64_empty() {
assert!((dot_product_f64(&[], &[]) - 0.0).abs() < 1e-10);
}
#[test]
fn dot_product_f64_precision() {
let n = 10_000;
let a = vec![1.0f32; n];
let b = vec![1.0f32; n];
assert!((dot_product_f64(&a, &b) - n as f64).abs() < 1e-10);
}
#[test]
fn dot_product_f64_odd_length() {
let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
let b = vec![5.0f32, 4.0, 3.0, 2.0, 1.0];
// 5 + 8 + 9 + 8 + 5 = 35
assert!((dot_product_f64(&a, &b) - 35.0).abs() < 1e-10);
}
// -----------------------------------------------------------------
// axpy
// -----------------------------------------------------------------
#[test]
fn axpy_basic() {
let x = vec![1.0f32, 2.0, 3.0];
let mut y = vec![10.0f32, 20.0, 30.0];
axpy(2.0, &x, &mut y);
assert_eq!(y, vec![12.0, 24.0, 36.0]);
}
#[test]
fn axpy_negative_alpha() {
let x = vec![1.0f32, 1.0, 1.0];
let mut y = vec![5.0f32, 5.0, 5.0];
axpy(-3.0, &x, &mut y);
assert_eq!(y, vec![2.0, 2.0, 2.0]);
}
// -----------------------------------------------------------------
// norm2
// -----------------------------------------------------------------
#[test]
fn norm2_basic() {
let x = vec![3.0f32, 4.0];
assert!((norm2(&x) - 5.0).abs() < 1e-10);
}
#[test]
fn norm2_zero() {
assert!((norm2(&vec![0.0f32; 5]) - 0.0).abs() < 1e-10);
}
// -----------------------------------------------------------------
// CG solver: convergence on well-conditioned systems
// -----------------------------------------------------------------
#[test]
fn cg_identity_matrix() {
let n = 5;
let matrix = identity(n);
let rhs = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-10, 100, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
for i in 0..n {
assert!(
(result.solution[i] as f64 - rhs[i]).abs() < 1e-5,
"x[{i}] = {} != {}",
result.solution[i],
rhs[i],
);
}
// Identity should converge in at most 1 iteration.
assert!(result.iterations <= 1);
}
#[test]
fn cg_diagonal_matrix() {
let diag = vec![2.0, 3.0, 5.0, 7.0];
let matrix = diagonal_matrix(&diag);
let rhs = vec![4.0, 9.0, 25.0, 49.0];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-10, 100, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
let expected = [2.0, 3.0, 5.0, 7.0];
for i in 0..4 {
assert!(
(result.solution[i] as f64 - expected[i]).abs() < 1e-4,
"x[{i}] = {} != {}",
result.solution[i],
expected[i],
);
}
}
#[test]
fn cg_tridiagonal_small() {
let n = 10;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 200, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert!(
result.residual_norm < 1e-6,
"residual = {}",
result.residual_norm,
);
assert!(
result.iterations <= n,
"took {} iterations for n={}",
result.iterations,
n,
);
}
#[test]
fn cg_tridiagonal_large() {
let n = 500;
let matrix = tridiagonal_spd(n);
let rhs: Vec<f64> = (0..n).map(|i| (i as f64 + 1.0) / n as f64).collect();
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 2000, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert!(
result.residual_norm < 1e-5,
"residual = {}",
result.residual_norm,
);
}
// -----------------------------------------------------------------
// Preconditioning
// -----------------------------------------------------------------
#[test]
fn cg_preconditioned_converges_faster() {
let n = 100;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
let budget = default_budget();
let no_precond = ConjugateGradientSolver::new(1e-8, 500, false);
let with_precond = ConjugateGradientSolver::new(1e-8, 500, true);
let result_no = no_precond.solve(&matrix, &rhs, &budget).unwrap();
let result_yes = with_precond.solve(&matrix, &rhs, &budget).unwrap();
assert!(result_no.residual_norm < 1e-6);
assert!(result_yes.residual_norm < 1e-6);
assert!(
result_yes.iterations <= result_no.iterations,
"preconditioned ({}) should use <= iterations than \
unpreconditioned ({})",
result_yes.iterations,
result_no.iterations,
);
}
// -----------------------------------------------------------------
// Zero RHS / empty system
// -----------------------------------------------------------------
#[test]
fn cg_zero_rhs() {
let matrix = tridiagonal_spd(5);
let rhs = vec![0.0f64; 5];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 100, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert_eq!(result.iterations, 0);
for &v in &result.solution {
assert!((v as f64).abs() < 1e-10);
}
}
#[test]
fn cg_empty_system() {
let matrix = CsrMatrix {
row_ptr: vec![0],
col_indices: vec![],
values: Vec::<f64>::new(),
rows: 0,
cols: 0,
};
let rhs: Vec<f64> = vec![];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 100, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert_eq!(result.iterations, 0);
assert!(result.solution.is_empty());
}
// -----------------------------------------------------------------
// Error cases
// -----------------------------------------------------------------
#[test]
fn cg_dimension_mismatch() {
let matrix = tridiagonal_spd(3);
let rhs = vec![1.0f64; 5];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 100, false);
let err = solver.solve(&matrix, &rhs, &budget).unwrap_err();
assert!(matches!(err, SolverError::InvalidInput(_)));
}
#[test]
fn cg_non_square_matrix() {
let matrix = CsrMatrix {
row_ptr: vec![0, 1, 2],
col_indices: vec![0, 1],
values: vec![1.0f64, 1.0],
rows: 2,
cols: 3,
};
let rhs = vec![1.0f64; 2];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 100, false);
let err = solver.solve(&matrix, &rhs, &budget).unwrap_err();
assert!(matches!(err, SolverError::InvalidInput(_)));
}
#[test]
fn cg_non_convergence() {
let n = 50;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
let budget = ComputeBudget {
max_time: Duration::from_secs(30),
max_iterations: 1,
tolerance: 1e-15,
};
let solver = ConjugateGradientSolver::new(1e-15, 1, false);
let err = solver.solve(&matrix, &rhs, &budget).unwrap_err();
assert!(matches!(err, SolverError::NonConvergence { .. }));
}
#[test]
fn cg_budget_iteration_limit() {
let n = 50;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
// Solver allows 1000, but budget allows only 2.
let solver = ConjugateGradientSolver::new(1e-15, 1000, false);
let budget = ComputeBudget {
max_time: Duration::from_secs(60),
max_iterations: 2,
tolerance: 1e-15,
};
let err = solver.solve(&matrix, &rhs, &budget).unwrap_err();
assert!(matches!(err, SolverError::NonConvergence { .. }));
}
// -----------------------------------------------------------------
// Convergence history
// -----------------------------------------------------------------
#[test]
fn cg_convergence_history_populated() {
let n = 20;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-10, 200, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert!(!result.convergence_history.is_empty());
// Final history entry should match the reported residual.
let last = result.convergence_history.last().unwrap();
assert!((last.residual_norm - result.residual_norm).abs() < 1e-12);
}
#[test]
fn cg_algorithm_field() {
let matrix = identity(3);
let rhs = vec![1.0f64; 3];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-8, 100, false);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
assert_eq!(result.algorithm, Algorithm::CG);
}
// -----------------------------------------------------------------
// Verify Ax = b for computed solution
// -----------------------------------------------------------------
#[test]
fn cg_solution_satisfies_system() {
let n = 20;
let matrix = tridiagonal_spd(n);
let rhs = vec![1.0f64; n];
let budget = default_budget();
let solver = ConjugateGradientSolver::new(1e-10, 200, true);
let result = solver.solve(&matrix, &rhs, &budget).unwrap();
// Up-cast solution back to f64 and compute Ax.
let x_f64: Vec<f64> = result.solution.iter().map(|&v| v as f64).collect();
let mut ax = vec![0.0f64; n];
matrix.spmv(&x_f64, &mut ax);
let mut max_err: f64 = 0.0;
for i in 0..n {
let err = (ax[i] - rhs[i]).abs();
if err > max_err {
max_err = err;
}
}
assert!(
max_err < 1e-4,
"max |Ax - b| = {max_err:.6e}, expected < 1e-4",
);
}
// -----------------------------------------------------------------
// Estimate complexity
// -----------------------------------------------------------------
#[test]
fn estimate_complexity_returns_cg() {
let solver = ConjugateGradientSolver::new(1e-8, 500, true);
let profile = SparsityProfile {
rows: 100,
cols: 100,
nnz: 298,
density: 0.0298,
is_diag_dominant: true,
estimated_spectral_radius: 0.5,
estimated_condition: 100.0,
is_symmetric_structure: true,
avg_nnz_per_row: 2.98,
max_nnz_per_row: 3,
};
let est = solver.estimate_complexity(&profile, 100);
assert_eq!(est.algorithm, Algorithm::CG);
assert_eq!(est.complexity_class, ComplexityClass::SqrtCondition);
assert!(est.estimated_iterations > 0);
assert!(est.estimated_flops > 0);
assert!(est.estimated_memory_bytes > 0);
}
// -----------------------------------------------------------------
// Accessors
// -----------------------------------------------------------------
#[test]
fn accessors() {
let solver = ConjugateGradientSolver::new(1e-6, 500, true);
assert!((solver.tolerance() - 1e-6).abs() < 1e-15);
assert_eq!(solver.max_iterations(), 500);
assert!(solver.use_preconditioner());
}
}
Reference in New Issue View Git Blame Copy Permalink