axiolid_numeric/
sparse.rs

1//! Sparse matrices and the preconditioned conjugate-gradient solve.
2//!
3//! Mesh operators (graph and cotangent Laplacians, their squares, mass
4//! matrices) have a handful of non-zeros per row, so a dense [`Matrix`]
5//! of a few thousand mesh vertices already costs more memory than the mesh
6//! itself and an `O(n^3)` factorisation more time than the rest of the
7//! algorithm. [`SparseMatrix`] stores only the non-zeros in compressed rows,
8//! and [`conjugate_gradient`] solves a symmetric positive definite system
9//! with it in `O(iterations * non-zeros)`.
10//!
11//! The solve follows Hestenes and Stiefel (1952) with a Jacobi (diagonal)
12//! preconditioner, the textbook form in Saad, *Iterative Methods for Sparse
13//! Linear Systems*, algorithm 9.1. Its stopping test is on the *true*
14//! residual `b - A x`, recomputed from the matrix when the recurrence says
15//! the tolerance is met, so drift in the recurred residual cannot report a
16//! convergence that did not happen.
17//!
18//! [`Matrix`]: crate::Matrix
19
20use crate::error::{all_finite, tolerance, NumericError, NumericResult, Status};
21
22/// A sparse matrix in compressed sparse row form.
23///
24/// Rows are stored in order and the column indices within each row are
25/// strictly increasing, so two matrices built from the same entries in any
26/// order are equal and every product is computed in the same order.
27#[derive(Debug, Clone, PartialEq)]
28pub struct SparseMatrix {
29    rows: usize,
30    cols: usize,
31    /// `row_start[i]..row_start[i + 1]` indexes row `i` in `columns` and `values`.
32    row_start: Vec<usize>,
33    columns: Vec<usize>,
34    values: Vec<f64>,
35}
36
37impl SparseMatrix {
38    /// Build from `(row, column, value)` triplets. Entries naming the same
39    /// position are summed, in the order given; entries that sum to zero are
40    /// kept as explicit zeros.
41    ///
42    /// # Errors
43    ///
44    /// [`NumericError::InvalidArgument`] for a triplet outside the
45    /// `rows x cols` shape and [`NumericError::NonFiniteInput`] for a
46    /// non-finite value.
47    pub fn from_triplets(
48        rows: usize,
49        cols: usize,
50        triplets: &[(usize, usize, f64)],
51    ) -> NumericResult<Self> {
52        let mut order: Vec<usize> = (0..triplets.len()).collect();
53        for &(i, j, v) in triplets {
54            if i >= rows || j >= cols {
55                return Err(NumericError::InvalidArgument {
56                    name: "triplet",
57                    reason: "must lie inside the matrix shape",
58                });
59            }
60            if !v.is_finite() {
61                return Err(NumericError::NonFiniteInput {
62                    name: "triplet value",
63                });
64            }
65        }
66        // Stable: equal positions keep their input order, so the summation
67        // order (and the rounding) is the caller's.
68        order.sort_by_key(|&k| (triplets[k].0, triplets[k].1));
69        let mut row_start = vec![0usize; rows + 1];
70        let mut columns: Vec<usize> = Vec::with_capacity(triplets.len());
71        let mut values: Vec<f64> = Vec::with_capacity(triplets.len());
72        let mut last: Option<(usize, usize)> = None;
73        for k in order {
74            let (i, j, v) = triplets[k];
75            if last == Some((i, j)) {
76                *values.last_mut().expect("an entry precedes its duplicate") += v;
77            } else {
78                columns.push(j);
79                values.push(v);
80                row_start[i + 1] += 1;
81                last = Some((i, j));
82            }
83        }
84        for i in 0..rows {
85            row_start[i + 1] += row_start[i];
86        }
87        Ok(Self {
88            rows,
89            cols,
90            row_start,
91            columns,
92            values,
93        })
94    }
95
96    /// Number of rows.
97    #[must_use]
98    pub fn rows(&self) -> usize {
99        self.rows
100    }
101
102    /// Number of columns.
103    #[must_use]
104    pub fn cols(&self) -> usize {
105        self.cols
106    }
107
108    /// Number of stored entries.
109    #[must_use]
110    pub fn nnz(&self) -> usize {
111        self.values.len()
112    }
113
114    /// Entry `(i, j)`; zero when it is not stored.
115    ///
116    /// # Panics
117    ///
118    /// When `(i, j)` is outside the shape.
119    #[must_use]
120    pub fn get(&self, i: usize, j: usize) -> f64 {
121        assert!(i < self.rows && j < self.cols, "matrix index out of range");
122        let (cols, vals) = self.row(i);
123        cols.binary_search(&j).map_or(0.0, |k| vals[k])
124    }
125
126    /// Column indices (increasing) and values of row `i`.
127    ///
128    /// # Panics
129    ///
130    /// When `i` is not a row.
131    #[must_use]
132    pub fn row(&self, i: usize) -> (&[usize], &[f64]) {
133        let range = self.row_start[i]..self.row_start[i + 1];
134        (&self.columns[range.clone()], &self.values[range])
135    }
136
137    /// `A x`.
138    ///
139    /// # Errors
140    ///
141    /// Refuses `x` of the wrong length.
142    pub fn mul_vec(&self, x: &[f64]) -> NumericResult<Vec<f64>> {
143        if x.len() != self.cols {
144            return Err(NumericError::DimensionMismatch {
145                name: "vector",
146                expected: self.cols,
147                found: x.len(),
148            });
149        }
150        let mut y = vec![0.0; self.rows];
151        self.mul_into(x, &mut y);
152        Ok(y)
153    }
154
155    fn mul_into(&self, x: &[f64], y: &mut [f64]) {
156        for (i, out) in y.iter_mut().enumerate() {
157            let (cols, vals) = self.row(i);
158            *out = cols.iter().zip(vals).map(|(&j, &v)| v * x[j]).sum();
159        }
160    }
161
162    /// Refuse a matrix that is not square or not symmetric to
163    /// `16 eps max(|a_ij|, |a_ji|)`.
164    fn require_symmetric(&self) -> NumericResult<()> {
165        if self.rows != self.cols {
166            return Err(NumericError::DimensionMismatch {
167                name: "symmetric matrix",
168                expected: self.rows,
169                found: self.cols,
170            });
171        }
172        for i in 0..self.rows {
173            let (cols, vals) = self.row(i);
174            for (&j, &v) in cols.iter().zip(vals) {
175                if j <= i {
176                    continue;
177                }
178                let w = self.get(j, i);
179                if (v - w).abs() > 16.0 * f64::EPSILON * v.abs().max(w.abs()) {
180                    return Err(NumericError::NotSymmetric { row: i, column: j });
181                }
182            }
183        }
184        Ok(())
185    }
186}
187
188/// Stopping rule and budget for [`conjugate_gradient`].
189#[derive(Debug, Clone, Copy, PartialEq)]
190pub struct ConjugateGradientOptions {
191    /// Stop when `||b - A x||_2 <= relative_tolerance * ||b||_2`.
192    pub relative_tolerance: f64,
193    /// Iteration budget; `None` allows `10 n + 100` for an `n x n` system.
194    pub max_iterations: Option<usize>,
195}
196
197impl Default for ConjugateGradientOptions {
198    fn default() -> Self {
199        Self {
200            relative_tolerance: 1e-12,
201            max_iterations: None,
202        }
203    }
204}
205
206/// The result of [`conjugate_gradient`].
207#[derive(Debug, Clone, PartialEq)]
208#[must_use]
209pub struct ConjugateGradientSolution {
210    /// The final iterate.
211    pub x: Vec<f64>,
212    /// `||b - A x||_2 / ||b||_2`, recomputed from the matrix at the final
213    /// iterate (zero when `b` is zero). A computed value, not an estimate.
214    pub relative_residual: f64,
215    /// Iterations taken.
216    pub iterations: usize,
217    /// [`Status::Converged`] when the relative residual met the tolerance,
218    /// [`Status::BudgetExhausted`] when the iteration budget ran out first.
219    pub status: Status,
220}
221
222impl ConjugateGradientSolution {
223    /// The solution when its tolerance was met.
224    ///
225    /// # Errors
226    ///
227    /// [`NumericError::NotConverged`] with the relative residual reached.
228    pub fn converged(self) -> NumericResult<Self> {
229        if self.status == Status::Converged {
230            Ok(self)
231        } else {
232            Err(NumericError::NotConverged {
233                name: "conjugate gradient",
234                error_estimate: self.relative_residual,
235            })
236        }
237    }
238}
239
240fn dot(a: &[f64], b: &[f64]) -> f64 {
241    a.iter().zip(b).map(|(x, y)| x * y).sum()
242}
243
244fn norm2(a: &[f64]) -> f64 {
245    dot(a, a).sqrt()
246}
247
248/// Solve `A x = b` for a symmetric positive definite sparse `A` by the
249/// Jacobi-preconditioned conjugate-gradient method, starting from
250/// `initial` (zero when `None`).
251///
252/// A run that exhausts its budget still returns its last iterate and
253/// residual with [`Status::BudgetExhausted`]; call
254/// [`ConjugateGradientSolution::converged`] to refuse it. A zero `b` returns
255/// the exact zero solution without iterating.
256///
257/// # Errors
258///
259/// - [`NumericError::DimensionMismatch`] for a non-square matrix or a
260///   wrongly sized `b` or `initial`.
261/// - [`NumericError::NonFiniteInput`] for a non-finite `b` or `initial`,
262///   [`NumericError::InvalidArgument`] for a negative or non-finite
263///   tolerance.
264/// - [`NumericError::NotSymmetric`] for an asymmetric matrix.
265/// - [`NumericError::NotPositiveDefinite`] when a diagonal entry is not
266///   positive (`pivot` is its row) or when the iteration meets a direction
267///   `p` with `p^T A p <= 0`, which proves `A` is not positive definite
268///   (`pivot` is the iteration).
269pub fn conjugate_gradient(
270    a: &SparseMatrix,
271    b: &[f64],
272    initial: Option<&[f64]>,
273    options: ConjugateGradientOptions,
274) -> NumericResult<ConjugateGradientSolution> {
275    a.require_symmetric()?;
276    let n = a.rows;
277    if b.len() != n {
278        return Err(NumericError::DimensionMismatch {
279            name: "right-hand side",
280            expected: n,
281            found: b.len(),
282        });
283    }
284    all_finite(b, "right-hand side")?;
285    let relative_tolerance = tolerance(options.relative_tolerance, "relative tolerance")?;
286    let mut x = match initial {
287        Some(x0) => {
288            if x0.len() != n {
289                return Err(NumericError::DimensionMismatch {
290                    name: "initial guess",
291                    expected: n,
292                    found: x0.len(),
293                });
294            }
295            all_finite(x0, "initial guess")?;
296            x0.to_vec()
297        }
298        None => vec![0.0; n],
299    };
300    let mut inverse_diagonal = Vec::with_capacity(n);
301    for i in 0..n {
302        let d = a.get(i, i);
303        if d <= 0.0 {
304            return Err(NumericError::NotPositiveDefinite { pivot: i });
305        }
306        inverse_diagonal.push(1.0 / d);
307    }
308    let b_norm = norm2(b);
309    if b_norm == 0.0 {
310        return Ok(ConjugateGradientSolution {
311            x: vec![0.0; n],
312            relative_residual: 0.0,
313            iterations: 0,
314            status: Status::Converged,
315        });
316    }
317    let target = relative_tolerance * b_norm;
318    let budget = options.max_iterations.unwrap_or(10 * n + 100);
319
320    let mut ax = vec![0.0; n];
321    let true_residual = |x: &[f64], ax: &mut Vec<f64>, r: &mut Vec<f64>| {
322        a.mul_into(x, ax);
323        for i in 0..n {
324            r[i] = b[i] - ax[i];
325        }
326        norm2(r)
327    };
328    let mut r = vec![0.0; n];
329    let mut r_norm = true_residual(&x, &mut ax, &mut r);
330    let mut z: Vec<f64> = r
331        .iter()
332        .zip(&inverse_diagonal)
333        .map(|(r, d)| r * d)
334        .collect();
335    let mut p = z.clone();
336    let mut rz = dot(&r, &z);
337    let mut ap = vec![0.0; n];
338    let mut iterations = 0;
339    loop {
340        if r_norm <= target {
341            // The recurred residual says done; confirm on the true one and
342            // restart from it if rounding let the two drift apart.
343            r_norm = true_residual(&x, &mut ax, &mut r);
344            if r_norm <= target {
345                return Ok(ConjugateGradientSolution {
346                    x,
347                    relative_residual: r_norm / b_norm,
348                    iterations,
349                    status: Status::Converged,
350                });
351            }
352            for i in 0..n {
353                z[i] = r[i] * inverse_diagonal[i];
354            }
355            p.clone_from(&z);
356            rz = dot(&r, &z);
357        }
358        if iterations >= budget {
359            r_norm = true_residual(&x, &mut ax, &mut r);
360            let status = if r_norm <= target {
361                Status::Converged
362            } else {
363                Status::BudgetExhausted
364            };
365            return Ok(ConjugateGradientSolution {
366                x,
367                relative_residual: r_norm / b_norm,
368                iterations,
369                status,
370            });
371        }
372        a.mul_into(&p, &mut ap);
373        let curvature = dot(&p, &ap);
374        if curvature <= 0.0 {
375            return Err(NumericError::NotPositiveDefinite { pivot: iterations });
376        }
377        let alpha = rz / curvature;
378        for i in 0..n {
379            x[i] += alpha * p[i];
380            r[i] -= alpha * ap[i];
381        }
382        r_norm = norm2(&r);
383        for i in 0..n {
384            z[i] = r[i] * inverse_diagonal[i];
385        }
386        let rz_next = dot(&r, &z);
387        let beta = rz_next / rz;
388        rz = rz_next;
389        for i in 0..n {
390            p[i] = z[i] + beta * p[i];
391        }
392        iterations += 1;
393    }
394}
axiolid_numeric/sparse.rs

axiolid_numeric/
sparse.rs
