Conjugate gradient

Deep dive

How it works

In: SPD operator A (matrix or matvec closure), b, preconditioner M⁻¹ (optional), tol · Out: x with ‖b − Ax‖ ≤ tol·‖b‖ · State: x, r, p and the scalar r·z (within one solve; none across solves) · Cost: one A·p per iteration; about ½·√κ·ln(2/ε) iterations

Preconditioned CG (x₀ = 0, r₀ = b):

z = M⁻¹ r;  p = z;  rz = r·z
repeat:
    q = A p
    α = rz / (p·q)
    x += α p;   r −= α q
    if ‖r‖ ≤ tol·‖b‖: stop
    z = M⁻¹ r
    β = (r·z)_new / rz;   rz = (r·z)_new
    p = z + β p

Residuals are mutually M⁻¹-orthogonal and directions are A-conjugate, so each step minimizes the A-norm error over a growing Krylov space span{z₀, M⁻¹Az₀, …}. That optimality is why CG beats steepest descent (which needs O(κ) iterations instead of O(√κ)).

Preconditioners, cheapest first:

• Jacobi M = diag(A): fixes bad row scaling (varying coefficients, mixed units); no help on a constant-coefficient Laplacian.
• Incomplete Cholesky IC(0): Cholesky restricted to A's sparsity pattern; usually fewer iterations than Jacobi (κ still O(h⁻²) for Poisson, with a smaller constant), but the triangular solves are sequential.
• Algebraic or geometric multigrid: one V-cycle as M⁻¹; iteration count nearly independent of mesh size.

Code

pub struct Csr { pub n: usize, pub row_ptr: Vec<usize>, pub col: Vec<u32>, pub val: Vec<f64> }

impl Csr {
    pub fn spmv(&self, x: &[f64], y: &mut [f64]) {
        for i in 0..self.n {
            let (s, e) = (self.row_ptr[i], self.row_ptr[i + 1]);
            y[i] = self.val[s..e].iter().zip(&self.col[s..e]).map(|(v, &c)| v * x[c as usize]).sum();
        }
    }
    pub fn diagonal(&self) -> Vec<f64> {
        (0..self.n).map(|i| {
            (self.row_ptr[i]..self.row_ptr[i + 1]).find(|&k| self.col[k] as usize == i).map_or(0.0, |k| self.val[k])
        }).collect()
    }
}

pub trait Precond { fn apply(&self, r: &[f64], z: &mut [f64]); }

pub struct Identity;
impl Precond for Identity { fn apply(&self, r: &[f64], z: &mut [f64]) { z.copy_from_slice(r); } }

pub struct Jacobi { inv_diag: Vec<f64> }
impl Jacobi {
    pub fn new(a: &Csr) -> Self {
        Jacobi { inv_diag: a.diagonal().iter().map(|&d| { assert!(d > 0.0, "SPD needs a positive diagonal"); 1.0 / d }).collect() }
    }
}
impl Precond for Jacobi {
    fn apply(&self, r: &[f64], z: &mut [f64]) { for i in 0..r.len() { z[i] = r[i] * self.inv_diag[i]; } }
}

fn dot(a: &[f64], b: &[f64]) -> f64 { a.iter().zip(b).map(|(x, y)| x * y).sum() }

pub struct CgResult { pub iterations: usize, pub rel_residual: f64, pub converged: bool }

/// Preconditioned CG. `x` holds the initial guess on entry and the solution on exit.
pub fn pcg(a: &Csr, b: &[f64], x: &mut [f64], m: &impl Precond, rtol: f64, max_iter: usize) -> CgResult {
    let n = b.len();
    let (mut r, mut z, mut q) = (vec![0.0; n], vec![0.0; n], vec![0.0; n]);
    a.spmv(x, &mut q);
    for i in 0..n { r[i] = b[i] - q[i]; }
    let bnorm = dot(b, b).sqrt().max(f64::MIN_POSITIVE);
    let mut rnorm = dot(&r, &r).sqrt();
    if rnorm <= rtol * bnorm { return CgResult { iterations: 0, rel_residual: rnorm / bnorm, converged: true }; }
    m.apply(&r, &mut z);
    let mut p = z.clone();
    let mut rz = dot(&r, &z);
    for it in 1..=max_iter {
        a.spmv(&p, &mut q);
        let pq = dot(&p, &q);
        if pq <= 0.0 { return CgResult { iterations: it, rel_residual: rnorm / bnorm, converged: false }; } // not SPD
        let alpha = rz / pq;
        for i in 0..n { x[i] += alpha * p[i]; r[i] -= alpha * q[i]; }
        rnorm = dot(&r, &r).sqrt();
        if rnorm <= rtol * bnorm { return CgResult { iterations: it, rel_residual: rnorm / bnorm, converged: true }; }
        m.apply(&r, &mut z);
        let rz_new = dot(&r, &z);
        let beta = rz_new / rz;
        rz = rz_new;
        for i in 0..n { p[i] = z[i] + beta * p[i]; }
    }
    CgResult { iterations: max_iter, rel_residual: rnorm / bnorm, converged: false }
}

/// S L S with L the 5-point Laplacian on an m x m interior grid and S = diag(s(node)):
/// SPD for any positive s, badly row-scaled when s varies (test case for Jacobi).
pub fn scaled_laplacian(m: usize, s: impl Fn(usize) -> f64) -> Csr {
    let (mut row_ptr, mut col, mut val) = (vec![0usize], Vec::new(), Vec::new());
    for j in 0..m {
        for i in 0..m {
            let sr = s(j * m + i);
            let mut push = |c: usize, v: f64| { col.push(c as u32); val.push(v * sr * s(c)); };
            if j > 0 { push((j - 1) * m + i, -1.0); }
            if i > 0 { push(j * m + i - 1, -1.0); }
            push(j * m + i, 4.0);
            if i + 1 < m { push(j * m + i + 1, -1.0); }
            if j + 1 < m { push((j + 1) * m + i, -1.0); }
            row_ptr.push(col.len());
        }
    }
    Csr { n: m * m, row_ptr, col, val }
}

Complexity and constants

• Per iteration: one SpMV (≈ 2·nnz flops, memory-bound), 2 dot products, 3 vector updates; about 5 vectors of n doubles.
• Measured with the code above, 2D Laplacian, rtol 10⁻⁸, b = 1: 59, 119, 239, 470 iterations for m = 32, 64, 128, 256 (doubling with m, as √κ predicts). With rows rescaled by factors from 0.01 to 100 (S·L·S), plain CG needs 2,260-28,134 iterations; Jacobi brings it back to 120-860. On the unscaled Laplacian Jacobi gives exactly the same counts as no preconditioner.
• Theory: iterations for an ε reduction ≈ ½·√κ·ln(2/ε); for the 2D Laplacian κ ≈ 4/(π²h²) ≈ 0.4·(m+1)², so √κ ≈ 0.64·m.
• In parallel, each dot product is a global reduction (latency-bound at scale); pipelined CG variants overlap them with the SpMV.

Efficiency

• Cost: per iteration one SpMV (about 2·nnz flops, memory-bound: a CSR matrix streams order of 12 bytes per nonzero), 2 dot products and 3 vector updates; about 5 vectors of n doubles. Measured with the code above on the 2D Laplacian (rtol 10⁻⁸): 59, 119, 239, 470 iterations for m = 32, 64, 128, 256. Bound: memory bandwidth; global reductions (latency) at scale.
• Levers:
  • preconditioner matched to the problem — Jacobi cut 2,260-28,134 iterations to 120-860 on a badly row-scaled Laplacian (measured above) — varying coefficients, mixed units
  • multigrid as M⁻¹ — iteration count nearly independent of mesh size (multigrid page) — elliptic PDEs
  • warm start from the previous solution — 309 instead of 364 iterations (measured above) — time stepping, Newton inner solves
  • fused vector kernels — one pass over r, z, p instead of several — large n, bandwidth-bound
  • pipelined or s-step CG — fewer global synchronizations — many nodes
• Naive → standard → state of the art:
  • steepest descent · O(κ) iterations
  • plain CG · O(√κ) iterations (470 at m = 256 above)
  • multigrid-preconditioned CG with fused or pipelined kernels · iteration count about constant in mesh size

Implementation notes

• Stopping test: relative residual ‖r‖/‖b‖ ≤ rtol; for inner solves of Newton use an Eisenstat-Walker forcing term instead of a fixed tolerance.
• Recurrence drift: the updated r drifts from b − Ax over thousands of iterations; recompute the true residual at the end (or periodically) before declaring convergence.
• Breakdown check: pᵀAp ≤ 0 means A is not SPD (or rounding on a nearly singular A); return instead of producing NaN.
• Warm start: in time-stepping, start from the previous solution; with a relative tolerance the saving is the log of how close it is (measured: b changed by 1 %, 309 instead of 364 iterations at m = 128).
• Symmetric preconditioner required: Gauss-Seidel must be symmetric (forward then backward sweep, SGS/SSOR) to be a valid PCG preconditioner.
• Fusing kernels: compute r·r and r·z in the same pass as the update; memory traffic, not flops, dominates.
• Singular A (pure Neumann Poisson, graph Laplacian): project b onto range(A) (subtract its mean) and project iterates, or CG drifts in the null space.

Pitfalls

• Using CG on a nonsymmetric matrix (upwinded advection, nonsymmetric Dirichlet elimination): it may stall or diverge silently.
• A preconditioner that is not SPD or changes between iterations: use flexible CG (FCG) if it varies (e.g. inexact inner solves).
• Absolute tolerances on badly scaled problems: either stops immediately or never; use relative.
• Expecting Jacobi to help a constant-coefficient Poisson problem: it only rescales by a constant.
• Not counting the cost of the preconditioner: IC(0) may halve iterations but double time per iteration and not parallelize.

Variants and alternatives

In production

Scientific computing: PETSc KSPCG, hypre (with BoomerAMG preconditioning), Trilinos Belos, SciPy scipy.sparse.linalg.cg, Eigen ConjugateGradient, cuSPARSE/AmgX on GPUs, the HPCG benchmark. Graphics: pressure solvers in fluid simulation (PCG with modified incomplete Cholesky, as in Bridson's book). Machine learning: GPyTorch performs most Gaussian-process inference with preconditioned CG; TRPO solves for the natural-gradient direction with CG.

Implementations

Sources

• Magnus R. Hestenes, Eduard Stiefel, Methods of Conjugate Gradients for Solving Linear Systems, Journal of Research of the National Bureau of Standards 49(6), 1952. paper
• Jonathan R. Shewchuk, An Introduction to the Conjugate Gradient Method Without the Agonizing Pain, Carnegie Mellon University technical report, 1994. pdf
• Yousef Saad, Iterative Methods for Sparse Linear Systems, 2nd ed., SIAM, 2003. book
• Anne Greenbaum, Iterative Methods for Solving Linear Systems, SIAM, 1997. book
• Robert Bridson, Fluid Simulation for Computer Graphics, 2nd ed., CRC Press, 2015. book
• Jacob R. Gardner et al., GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration, NeurIPS 2018. paper · code
• John Schulman et al., Trust Region Policy Optimization, ICML 2015. paper