Skip to main content

oxicuda_sparse/preconditioner/
amg.rs

1//! Algebraic multigrid via smoothed aggregation (Vaněk-Mandel-Brezina).
2//!
3//! This module builds a multigrid hierarchy for a symmetric, (near) positive-
4//! definite sparse matrix `A` purely from its algebraic entries -- no geometric
5//! grid information is required. The setup phase coarsens `A` by aggregating
6//! strongly-connected nodes, builds a tentative piecewise-constant interpolation
7//! that exactly represents the near-null-space (the constant vector), smooths it
8//! with one Jacobi step to improve interpolation quality, and forms the coarse
9//! operator through the Galerkin triple product `A_c = Pᵀ A P`. The solve phase
10//! applies V-cycles: weighted-Jacobi pre/post smoothing on each level and a
11//! direct dense solve on the coarsest level.
12//!
13//! ## Setup (smoothed aggregation)
14//!
15//! 1. **Strength of connection.** Nodes `i` and `j` are strongly connected when
16//!    `|A_ij| ≥ θ · sqrt(|A_ii| · |A_jj|)`, with `θ = strength_threshold`.
17//! 2. **Aggregation (greedy two-pass).** Pass 1 forms a new aggregate
18//!    `{i} ∪ strong_neighbours(i)` for every node `i` whose strong neighbours
19//!    are all still free. Pass 2 attaches each leftover node to an adjacent
20//!    existing aggregate, or makes it a singleton.
21//! 3. **Tentative prolongator `P0`.** Column `g` is the constant near-null-space
22//!    vector restricted to aggregate `g`, normalised to unit Euclidean norm, so
23//!    `range(P0)` contains the constant and `P0` has orthonormal columns.
24//! 4. **Smoothed prolongator.** `P = (I − ω D⁻¹ A) P0` with
25//!    `ω = smooth_omega / ρ̂(D⁻¹A)`, where `ρ̂` is estimated by power iteration.
26//! 5. **Galerkin coarse operator.** `A_c = Pᵀ A P`, computed with the host
27//!    sparse-sparse product.
28//!
29//! The recursion stops when the operator order drops to `max_coarse` or the
30//! level count reaches `max_levels`.
31
32use crate::error::{SparseError, SparseResult};
33use crate::host_csr::{HostCsr, dense_solve};
34
35/// A single level of the multigrid hierarchy.
36pub struct AmgLevel {
37    /// The (fine) operator on this level.
38    pub a: HostCsr,
39    /// The prolongation (interpolation) operator to this level from the next
40    /// coarser level: `fine = P · coarse`.
41    pub p: HostCsr,
42    /// The restriction operator `R = Pᵀ`.
43    pub r: HostCsr,
44    /// The diagonal of `a`, cached for weighted-Jacobi smoothing.
45    pub diag: Vec<f64>,
46}
47
48/// A complete smoothed-aggregation multigrid hierarchy.
49pub struct AmgHierarchy {
50    /// Levels ordered finest-first. Each level carries its operator and the
51    /// transfer operators connecting it to the next coarser level.
52    pub levels: Vec<AmgLevel>,
53    /// The coarsest operator, solved directly during the V-cycle.
54    pub coarse: HostCsr,
55    /// Number of weighted-Jacobi pre-smoothing sweeps applied per level.
56    pub pre_sweeps: usize,
57    /// Number of weighted-Jacobi post-smoothing sweeps applied per level.
58    pub post_sweeps: usize,
59}
60
61/// Tunable parameters for the smoothed-aggregation setup and the V-cycle.
62pub struct AmgOptions {
63    /// Strength-of-connection threshold `θ` (e.g. `0.08`).
64    pub strength_threshold: f64,
65    /// Maximum number of levels (including the finest) before forcing a coarse
66    /// solve.
67    pub max_levels: usize,
68    /// Stop coarsening once the operator order is `≤ max_coarse`.
69    pub max_coarse: usize,
70    /// Prolongator-smoothing weight numerator; the effective Jacobi weight is
71    /// `smooth_omega / ρ̂(D⁻¹A)` (e.g. `smooth_omega = 4/3`).
72    pub smooth_omega: f64,
73    /// Number of pre-smoothing sweeps per level.
74    pub pre_sweeps: usize,
75    /// Number of post-smoothing sweeps per level.
76    pub post_sweeps: usize,
77}
78
79impl Default for AmgOptions {
80    fn default() -> Self {
81        Self {
82            strength_threshold: 0.08,
83            max_levels: 25,
84            max_coarse: 16,
85            smooth_omega: 4.0 / 3.0,
86            pre_sweeps: 2,
87            post_sweeps: 2,
88        }
89    }
90}
91
92/// Builds a smoothed-aggregation multigrid hierarchy for `a`.
93///
94/// # Errors
95///
96/// Returns [`SparseError::DimensionMismatch`] if `a` is not square,
97/// [`SparseError::InvalidArgument`] if `a` is empty or an option is out of
98/// range, [`SparseError::SingularMatrix`] if a zero diagonal is encountered, or
99/// propagates errors from the Galerkin product.
100pub fn amg_setup(a: &HostCsr, opts: &AmgOptions) -> SparseResult<AmgHierarchy> {
101    if a.nrows != a.ncols {
102        return Err(SparseError::DimensionMismatch(format!(
103            "AMG requires a square matrix, got {}x{}",
104            a.nrows, a.ncols
105        )));
106    }
107    if a.nrows == 0 {
108        return Err(SparseError::InvalidArgument(
109            "cannot build a hierarchy for an empty matrix".to_string(),
110        ));
111    }
112    if opts.max_levels == 0 {
113        return Err(SparseError::InvalidArgument(
114            "max_levels must be >= 1".to_string(),
115        ));
116    }
117
118    let mut levels: Vec<AmgLevel> = Vec::new();
119    let mut current = a.clone();
120
121    loop {
122        let n = current.nrows;
123        // Stop if we are small enough or have reached the level cap. The number
124        // of stored levels is the number of coarsening steps; the finest level
125        // is level 0.
126        if n <= opts.max_coarse || levels.len() + 1 >= opts.max_levels {
127            break;
128        }
129
130        // 1. Strength of connection.
131        let strength = strength_of_connection(&current, opts.strength_threshold);
132
133        // 2. Aggregation.
134        let (aggregates, num_agg) = aggregate(&current, &strength);
135        // Guard against failure to coarsen (e.g. every node a singleton): if no
136        // reduction occurred, stop to avoid an infinite hierarchy.
137        if num_agg == 0 || num_agg >= n {
138            break;
139        }
140
141        // 3. Tentative prolongator P0 (column-normalised constant).
142        let p0 = tentative_prolongator(&aggregates, num_agg, n);
143
144        // 4. Smoothed prolongator P = (I - omega D^{-1} A) P0.
145        let diag = current.diagonal();
146        for &d in &diag {
147            if d == 0.0 {
148                return Err(SparseError::SingularMatrix);
149            }
150        }
151        let rho = spectral_radius_estimate(&current, &diag);
152        let omega = opts.smooth_omega / rho.max(1e-30);
153        let p = smooth_prolongator(&current, &diag, &p0, omega)?;
154
155        // 5. Galerkin coarse operator A_c = Pᵀ A P.
156        let r = p.transpose();
157        let ap = current.matmul(&p)?;
158        let a_coarse = r.matmul(&ap)?;
159
160        let level = AmgLevel {
161            a: current.clone(),
162            p,
163            r,
164            diag,
165        };
166        levels.push(level);
167
168        current = a_coarse;
169    }
170
171    Ok(AmgHierarchy {
172        levels,
173        coarse: current,
174        pre_sweeps: opts.pre_sweeps.max(1),
175        post_sweeps: opts.post_sweeps.max(1),
176    })
177}
178
179/// Applies one multigrid V-cycle, updating `x` in place toward the solution of
180/// `A x = b` on the finest level.
181///
182/// The vectors `b` and `x` must have length equal to the order of the finest
183/// operator. Out-of-range lengths are clamped by the smoother bounds and leave
184/// the unaffected entries untouched.
185pub fn amg_v_cycle(h: &AmgHierarchy, b: &[f64], x: &mut [f64]) {
186    v_cycle_recursive(h, 0, b, x);
187}
188
189/// Recursive V-cycle on level `level`. `b` is the right-hand side for the
190/// operator at `level`, `x` the current iterate (updated in place).
191fn v_cycle_recursive(h: &AmgHierarchy, level: usize, b: &[f64], x: &mut [f64]) {
192    if level >= h.levels.len() {
193        // Coarsest level: direct dense solve.
194        let n = h.coarse.nrows;
195        if n == 0 {
196            return;
197        }
198        let dense = h.coarse.to_dense();
199        match dense_solve(&dense, b, n) {
200            Ok(sol) => {
201                x[..n].copy_from_slice(&sol[..n]);
202            }
203            Err(_) => {
204                // Fall back to a few Jacobi sweeps if the coarse solve is
205                // singular (e.g. a pure-Neumann constant null space).
206                let diag = h.coarse.diagonal();
207                jacobi_sweeps(&h.coarse, &diag, b, x, 30, 2.0 / 3.0);
208            }
209        }
210        return;
211    }
212
213    let lvl = &h.levels[level];
214
215    // Effective smoothing weight: weighted Jacobi with omega = 2/3.
216    let omega = 2.0 / 3.0;
217
218    // Pre-smoothing.
219    jacobi_sweeps(&lvl.a, &lvl.diag, b, x, h.pre_sweeps, omega);
220
221    // Residual r = b - A x.
222    let ax = lvl.a.matvec(x);
223    let mut residual = vec![0.0f64; b.len()];
224    for i in 0..b.len() {
225        residual[i] = b[i] - ax[i];
226    }
227
228    // Restrict to the coarse level: r_c = R r.
229    let r_coarse = lvl.r.matvec(&residual);
230    let coarse_n = r_coarse.len();
231    let mut e_coarse = vec![0.0f64; coarse_n];
232
233    // Recurse to solve A_c e_c = r_c.
234    v_cycle_recursive(h, level + 1, &r_coarse, &mut e_coarse);
235
236    // Prolong the correction and apply: x += P e_c.
237    let correction = lvl.p.matvec(&e_coarse);
238    for i in 0..x.len().min(correction.len()) {
239        x[i] += correction[i];
240    }
241
242    // Post-smoothing.
243    jacobi_sweeps(&lvl.a, &lvl.diag, b, x, h.post_sweeps, omega);
244}
245
246/// Iteratively solves `A x = b` to a relative residual tolerance using V-cycles.
247///
248/// Returns `(x, iters, final_rel_residual)`.
249///
250/// # Errors
251///
252/// Propagates setup errors from [`amg_setup`]. If the right-hand side is the
253/// zero vector the solution is the zero vector and zero iterations are reported.
254pub fn amg_solve(
255    a: &HostCsr,
256    b: &[f64],
257    opts: &AmgOptions,
258    max_iter: usize,
259    tol: f64,
260) -> SparseResult<(Vec<f64>, usize, f64)> {
261    if b.len() != a.nrows {
262        return Err(SparseError::DimensionMismatch(format!(
263            "rhs length {} must equal matrix order {}",
264            b.len(),
265            a.nrows
266        )));
267    }
268    let hierarchy = amg_setup(a, opts)?;
269
270    let b_norm = norm(b);
271    if b_norm == 0.0 {
272        return Ok((vec![0.0; a.nrows], 0, 0.0));
273    }
274
275    let mut x = vec![0.0f64; a.nrows];
276    let mut final_rel = 1.0;
277    let mut iters = 0;
278    for it in 0..max_iter {
279        iters = it + 1;
280        amg_v_cycle(&hierarchy, b, &mut x);
281        let ax = a.matvec(&x);
282        let mut res = vec![0.0f64; b.len()];
283        for i in 0..b.len() {
284            res[i] = b[i] - ax[i];
285        }
286        final_rel = norm(&res) / b_norm;
287        if final_rel < tol {
288            break;
289        }
290    }
291
292    Ok((x, iters, final_rel))
293}
294
295// ---------------------------------------------------------------------------
296// Setup building blocks.
297// ---------------------------------------------------------------------------
298
299/// Per-node list of strongly-connected neighbours (excluding the node itself).
300type Strength = Vec<Vec<usize>>;
301
302/// Computes the strength-of-connection graph: `i ~ j` iff
303/// `|A_ij| ≥ θ · sqrt(|A_ii| · |A_jj|)`.
304fn strength_of_connection(a: &HostCsr, theta: f64) -> Strength {
305    let n = a.nrows;
306    let diag = a.diagonal();
307    let mut strength: Strength = vec![Vec::new(); n];
308    for i in 0..n {
309        let start = a.row_ptr[i];
310        let end = a.row_ptr[i + 1];
311        let dii = diag[i].abs();
312        for kk in start..end {
313            let j = a.col_indices[kk];
314            if j == i {
315                continue;
316            }
317            let aij = a.values[kk].abs();
318            let djj = diag[j].abs();
319            let bound = theta * (dii * djj).sqrt();
320            if aij >= bound && aij > 0.0 {
321                strength[i].push(j);
322            }
323        }
324    }
325    strength
326}
327
328/// Greedy two-pass aggregation. Returns `(aggregate_of_node, num_aggregates)`.
329fn aggregate(a: &HostCsr, strength: &Strength) -> (Vec<usize>, usize) {
330    let n = a.nrows;
331    const UNASSIGNED: usize = usize::MAX;
332    let mut agg = vec![UNASSIGNED; n];
333    let mut num_agg = 0usize;
334
335    // Pass 1: form root aggregates from nodes whose strong neighbourhood is
336    // entirely free.
337    for i in 0..n {
338        if agg[i] != UNASSIGNED {
339            continue;
340        }
341        let all_free = strength[i].iter().all(|&j| agg[j] == UNASSIGNED);
342        if !all_free {
343            continue;
344        }
345        // Create a new aggregate containing i and its strong neighbours.
346        let g = num_agg;
347        num_agg += 1;
348        agg[i] = g;
349        for &j in &strength[i] {
350            agg[j] = g;
351        }
352    }
353
354    // Pass 2: sweep remaining nodes, attaching each to the aggregate of a
355    // strong neighbour (preferring the most strongly-connected one). Nodes with
356    // no aggregated strong neighbour become singletons.
357    for i in 0..n {
358        if agg[i] != UNASSIGNED {
359            continue;
360        }
361        // Find the best already-aggregated strong neighbour.
362        let mut best: Option<usize> = None;
363        let mut best_mag = 0.0f64;
364        let start = a.row_ptr[i];
365        let end = a.row_ptr[i + 1];
366        for &j in &strength[i] {
367            if agg[j] != UNASSIGNED {
368                // Magnitude of A_ij for tie-breaking.
369                let mut mag = 0.0;
370                for kk in start..end {
371                    if a.col_indices[kk] == j {
372                        mag = a.values[kk].abs();
373                        break;
374                    }
375                }
376                if best.is_none() || mag > best_mag {
377                    best = Some(agg[j]);
378                    best_mag = mag;
379                }
380            }
381        }
382        match best {
383            Some(g) => agg[i] = g,
384            None => {
385                let g = num_agg;
386                num_agg += 1;
387                agg[i] = g;
388            }
389        }
390    }
391
392    (agg, num_agg)
393}
394
395/// Builds the tentative prolongator `P0` (`n × num_agg`) whose column `g` is the
396/// constant near-null-space restricted to aggregate `g`, normalised to unit
397/// Euclidean norm. With this scaling `P0` has orthonormal columns and its range
398/// contains the constant vector.
399fn tentative_prolongator(agg: &[usize], num_agg: usize, n: usize) -> HostCsr {
400    // Aggregate sizes for normalisation (constant column has norm sqrt(size)).
401    let mut sizes = vec![0usize; num_agg];
402    for &g in agg {
403        sizes[g] += 1;
404    }
405    let mut row_ptr = vec![0usize; n + 1];
406    let mut col_indices = Vec::with_capacity(n);
407    let mut values = Vec::with_capacity(n);
408    for i in 0..n {
409        let g = agg[i];
410        let norm = (sizes[g] as f64).sqrt();
411        col_indices.push(g);
412        values.push(1.0 / norm);
413        row_ptr[i + 1] = col_indices.len();
414    }
415    HostCsr {
416        nrows: n,
417        ncols: num_agg,
418        row_ptr,
419        col_indices,
420        values,
421    }
422}
423
424/// Forms the smoothed prolongator `P = (I − ω D⁻¹ A) P0`.
425///
426/// Computed as `P = P0 − ω D⁻¹ (A P0)` row by row, where `D⁻¹` scales row `i`
427/// by `1 / A_ii`.
428fn smooth_prolongator(
429    a: &HostCsr,
430    diag: &[f64],
431    p0: &HostCsr,
432    omega: f64,
433) -> SparseResult<HostCsr> {
434    // A * P0 (sparse-sparse).
435    let ap0 = a.matmul(p0)?;
436    let n = a.nrows;
437    let ncols = p0.ncols;
438
439    // Accumulate each row of P = P0 - omega * Dinv * (A P0) in a dense buffer.
440    let mut row_ptr = vec![0usize; n + 1];
441    let mut col_indices: Vec<usize> = Vec::new();
442    let mut values: Vec<f64> = Vec::new();
443
444    let mut accum = vec![0.0f64; ncols];
445    let mut touched: Vec<usize> = Vec::new();
446    let mut is_touched = vec![false; ncols];
447
448    for i in 0..n {
449        let dinv = 1.0 / diag[i];
450
451        // P0 contribution (row i of P0 has a single entry).
452        let p0_s = p0.row_ptr[i];
453        let p0_e = p0.row_ptr[i + 1];
454        for kk in p0_s..p0_e {
455            let c = p0.col_indices[kk];
456            if !is_touched[c] {
457                is_touched[c] = true;
458                touched.push(c);
459            }
460            accum[c] += p0.values[kk];
461        }
462
463        // -omega * Dinv * (A P0) contribution.
464        let ap_s = ap0.row_ptr[i];
465        let ap_e = ap0.row_ptr[i + 1];
466        for kk in ap_s..ap_e {
467            let c = ap0.col_indices[kk];
468            if !is_touched[c] {
469                is_touched[c] = true;
470                touched.push(c);
471            }
472            accum[c] -= omega * dinv * ap0.values[kk];
473        }
474
475        touched.sort_unstable();
476        for &c in &touched {
477            let v = accum[c];
478            if v != 0.0 {
479                col_indices.push(c);
480                values.push(v);
481            }
482            accum[c] = 0.0;
483            is_touched[c] = false;
484        }
485        touched.clear();
486        row_ptr[i + 1] = col_indices.len();
487    }
488
489    Ok(HostCsr {
490        nrows: n,
491        ncols,
492        row_ptr,
493        col_indices,
494        values,
495    })
496}
497
498/// Estimates the spectral radius `ρ(D⁻¹A)` by a few power iterations starting
499/// from a deterministic vector. `D⁻¹A` is applied implicitly.
500fn spectral_radius_estimate(a: &HostCsr, diag: &[f64]) -> f64 {
501    let n = a.nrows;
502    if n == 0 {
503        return 1.0;
504    }
505    // Deterministic, non-constant starting vector.
506    let mut v: Vec<f64> = (0..n).map(|i| 1.0 + (i % 7) as f64 * 0.3).collect();
507    normalize(&mut v);
508
509    let mut lambda = 1.0;
510    for _ in 0..15 {
511        // w = D^{-1} A v.
512        let av = a.matvec(&v);
513        let mut w = vec![0.0f64; n];
514        for i in 0..n {
515            w[i] = av[i] / diag[i];
516        }
517        let nrm = norm(&w);
518        if nrm < 1e-300 {
519            break;
520        }
521        // Rayleigh-style estimate.
522        lambda = dot(&v, &w);
523        for i in 0..n {
524            v[i] = w[i] / nrm;
525        }
526    }
527    // For an SPD M-matrix, D^{-1}A has positive spectrum; guard the estimate.
528    lambda.abs().max(1.0)
529}
530
531// ---------------------------------------------------------------------------
532// Smoother and vector helpers.
533// ---------------------------------------------------------------------------
534
535/// Applies `sweeps` weighted-Jacobi sweeps `x ← x + ω D⁻¹ (b − A x)` in place.
536fn jacobi_sweeps(a: &HostCsr, diag: &[f64], b: &[f64], x: &mut [f64], sweeps: usize, omega: f64) {
537    let n = a.nrows;
538    for _ in 0..sweeps {
539        let ax = a.matvec(x);
540        for i in 0..n {
541            if diag[i] != 0.0 {
542                x[i] += omega * (b[i] - ax[i]) / diag[i];
543            }
544        }
545    }
546}
547
548/// Euclidean inner product.
549fn dot(a: &[f64], b: &[f64]) -> f64 {
550    a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum()
551}
552
553/// Euclidean norm.
554fn norm(a: &[f64]) -> f64 {
555    dot(a, a).sqrt()
556}
557
558/// Scales a vector to unit Euclidean norm (no-op for the zero vector).
559fn normalize(v: &mut [f64]) {
560    let nrm = norm(v);
561    if nrm > 1e-300 {
562        let inv = 1.0 / nrm;
563        for x in v.iter_mut() {
564            *x *= inv;
565        }
566    }
567}
568
569#[cfg(test)]
570mod tests {
571    use super::*;
572    use crate::host_csr::{laplacian_1d, laplacian_2d};
573
574    #[test]
575    fn poisson_1d_converges() {
576        // 1D Poisson: recover a known solution in few V-cycles.
577        let n = 64;
578        let a = laplacian_1d(n);
579        let x_star: Vec<f64> = (0..n).map(|i| ((i as f64) * 0.1).sin() + 1.0).collect();
580        let b = a.matvec(&x_star);
581        let opts = AmgOptions::default();
582        let (x, iters, rel) = amg_solve(&a, &b, &opts, 60, 1e-8).expect("solve");
583        assert!(rel < 1e-8, "did not converge: rel = {rel}");
584        assert!(iters < 20, "took too many V-cycles: {iters}");
585        for i in 0..n {
586            assert!(
587                (x[i] - x_star[i]).abs() < 1e-6,
588                "solution mismatch at {i}: {} vs {}",
589                x[i],
590                x_star[i]
591            );
592        }
593    }
594
595    #[test]
596    fn poisson_2d_geometric_residual_drop() {
597        // 2D Poisson on a 16x16 grid: residual must drop geometrically.
598        let g = 16;
599        let a = laplacian_2d(g, g);
600        let n = g * g;
601        let x_star: Vec<f64> = (0..n).map(|i| 1.0 + (i % 5) as f64).collect();
602        let b = a.matvec(&x_star);
603        let opts = AmgOptions::default();
604        let hierarchy = amg_setup(&a, &opts).expect("setup");
605
606        let b_norm = norm(&b);
607        let mut x = vec![0.0f64; n];
608        let mut residuals = Vec::new();
609        for _ in 0..6 {
610            amg_v_cycle(&hierarchy, &b, &mut x);
611            let ax = a.matvec(&x);
612            let mut res = vec![0.0; n];
613            for i in 0..n {
614                res[i] = b[i] - ax[i];
615            }
616            residuals.push(norm(&res) / b_norm);
617        }
618        // Geometric drop: each cycle reduces the residual by at least 0.7x.
619        for w in residuals.windows(2) {
620            if w[0] > 1e-12 {
621                assert!(
622                    w[1] < 0.7 * w[0],
623                    "residual did not drop geometrically: {} -> {}",
624                    w[0],
625                    w[1]
626                );
627            }
628        }
629    }
630
631    #[test]
632    fn mesh_independence() {
633        // Iteration counts to reach tol stay bounded as n grows -- the
634        // hallmark of multigrid.
635        let opts = AmgOptions::default();
636        let tol = 1e-8;
637
638        let a64 = laplacian_1d(64);
639        let xs64: Vec<f64> = (0..64).map(|i| 1.0 + (i % 3) as f64).collect();
640        let b64 = a64.matvec(&xs64);
641        let (_x, iters64, _r) = amg_solve(&a64, &b64, &opts, 50, tol).expect("solve64");
642
643        let g = 16;
644        let a256 = laplacian_2d(g, g);
645        let xs256: Vec<f64> = (0..256).map(|i| 1.0 + (i % 4) as f64).collect();
646        let b256 = a256.matvec(&xs256);
647        let (_x2, iters256, _r2) = amg_solve(&a256, &b256, &opts, 50, tol).expect("solve256");
648
649        assert!(iters64 < 25, "n=64 took {iters64} cycles");
650        assert!(iters256 < 25, "n=256 took {iters256} cycles");
651    }
652
653    #[test]
654    fn galerkin_coarse_is_symmetric() {
655        let a = laplacian_2d(8, 8);
656        let opts = AmgOptions::default();
657        let h = amg_setup(&a, &opts).expect("setup");
658        // Check every coarse operator (including the final one) is symmetric.
659        for lvl in &h.levels {
660            let ac = if lvl.a.nrows == a.nrows {
661                // skip the finest; check the coarse via Galerkin on this level
662                let ap = lvl.a.matmul(&lvl.p).expect("ap");
663                lvl.r.matmul(&ap).expect("rap")
664            } else {
665                lvl.a.clone()
666            };
667            for i in 0..ac.nrows {
668                for j in 0..ac.ncols {
669                    let aij = ac.get(i, j).unwrap_or(0.0);
670                    let aji = ac.get(j, i).unwrap_or(0.0);
671                    assert!(
672                        (aij - aji).abs() < 1e-9,
673                        "coarse operator not symmetric at ({i},{j}): {aij} vs {aji}"
674                    );
675                }
676            }
677        }
678        // Final coarse operator symmetric too.
679        let ac = &h.coarse;
680        for i in 0..ac.nrows {
681            for j in 0..ac.ncols {
682                assert!((ac.get(i, j).unwrap_or(0.0) - ac.get(j, i).unwrap_or(0.0)).abs() < 1e-9);
683            }
684        }
685    }
686
687    #[test]
688    fn tentative_prolongator_preserves_constant() {
689        // P0 · 1_coarse == 1_fine exactly, where 1_coarse[g] = sqrt(size(g))
690        // because the columns are unit-normalised constants.
691        let a = laplacian_2d(6, 6);
692        let opts = AmgOptions::default();
693        let strength = strength_of_connection(&a, opts.strength_threshold);
694        let (agg, num_agg) = aggregate(&a, &strength);
695        let p0 = tentative_prolongator(&agg, num_agg, a.nrows);
696        // Build the coarse constant representation.
697        let mut sizes = vec![0usize; num_agg];
698        for &g in &agg {
699            sizes[g] += 1;
700        }
701        let one_coarse: Vec<f64> = sizes.iter().map(|&s| (s as f64).sqrt()).collect();
702        let one_fine = p0.matvec(&one_coarse);
703        for &v in &one_fine {
704            assert!(
705                (v - 1.0).abs() < 1e-12,
706                "P0 does not preserve constant: {v}"
707            );
708        }
709    }
710
711    #[test]
712    fn smoothed_prolongator_approximates_constant() {
713        // The smoothed P interpolates the constant up to the boundary smoothing
714        // term ω D⁻¹ A 1, which is zero in the interior of the Laplacian.
715        let a = laplacian_1d(32);
716        let opts = AmgOptions::default();
717        let strength = strength_of_connection(&a, opts.strength_threshold);
718        let (agg, num_agg) = aggregate(&a, &strength);
719        let p0 = tentative_prolongator(&agg, num_agg, a.nrows);
720        let diag = a.diagonal();
721        let rho = spectral_radius_estimate(&a, &diag);
722        let omega = opts.smooth_omega / rho;
723        let p = smooth_prolongator(&a, &diag, &p0, omega).expect("smooth");
724
725        let mut sizes = vec![0usize; num_agg];
726        for &g in &agg {
727            sizes[g] += 1;
728        }
729        let one_coarse: Vec<f64> = sizes.iter().map(|&s| (s as f64).sqrt()).collect();
730        let one_fine = p.matvec(&one_coarse);
731        // Interior nodes (away from the two boundaries) reproduce the constant
732        // exactly because A·1 = 0 there.
733        let mut interior_ok = 0;
734        for (i, &v) in one_fine.iter().enumerate() {
735            if i > 2 && i < a.nrows - 3 {
736                assert!(
737                    (v - 1.0).abs() < 1e-9,
738                    "smoothed P does not preserve constant in interior at {i}: {v}"
739                );
740                interior_ok += 1;
741            }
742        }
743        assert!(interior_ok > 0);
744    }
745
746    #[test]
747    fn rejects_non_square() {
748        let a = HostCsr::new(2, 3, vec![0, 1, 2], vec![0, 1], vec![1.0, 1.0]).expect("valid");
749        let opts = AmgOptions::default();
750        assert!(matches!(
751            amg_setup(&a, &opts),
752            Err(SparseError::DimensionMismatch(_))
753        ));
754    }
755
756    #[test]
757    fn strength_threshold_behaviour() {
758        // With a high threshold no connections are strong; with a low threshold
759        // the off-diagonals are strong.
760        let a = laplacian_1d(10);
761        let high = strength_of_connection(&a, 0.9);
762        // |A_ij| = 1, sqrt(|A_ii A_jj|) = 2, so threshold 0.9 -> bound 1.8 > 1.
763        assert!(high.iter().all(|nbrs| nbrs.is_empty()));
764        let low = strength_of_connection(&a, 0.4);
765        // bound 0.8 < 1 -> all off-diagonals strong.
766        assert!(low.iter().any(|nbrs| !nbrs.is_empty()));
767    }
768
769    #[test]
770    fn aggregation_covers_all_nodes() {
771        let a = laplacian_2d(7, 7);
772        let opts = AmgOptions::default();
773        let strength = strength_of_connection(&a, opts.strength_threshold);
774        let (agg, num_agg) = aggregate(&a, &strength);
775        assert_eq!(agg.len(), a.nrows);
776        assert!(num_agg > 0 && num_agg < a.nrows);
777        for &g in &agg {
778            assert!(g < num_agg);
779        }
780    }
781
782    #[test]
783    fn empty_rhs_zero_solution() {
784        let a = laplacian_1d(16);
785        let b = vec![0.0; 16];
786        let opts = AmgOptions::default();
787        let (x, iters, rel) = amg_solve(&a, &b, &opts, 10, 1e-8).expect("solve");
788        assert_eq!(iters, 0);
789        assert_eq!(rel, 0.0);
790        assert!(x.iter().all(|&v| v == 0.0));
791    }
792
793    /// Cross-feature integration test: on a single 1-D Poisson matrix exercise
794    /// all three algorithms implemented in this crate and check they agree with
795    /// each other and with analytic theory.
796    ///
797    /// * LOBPCG recovers the smallest eigenvalue `λ₁ = 2 − 2cos(π/(n+1))`.
798    /// * The complete IC(k) factor solves `A x = b` exactly.
799    /// * AMG solves the same system to tolerance.
800    /// * The IC(k) and AMG solutions match.
801    #[test]
802    fn cross_feature_poisson_consistency() {
803        use crate::eig::lobpcg::lobpcg;
804        use crate::preconditioner::ick::ic_k;
805        use std::f64::consts::PI;
806
807        let n = 48;
808        let a = laplacian_1d(n);
809
810        // (1) LOBPCG: smallest eigenvalue matches analytic value.
811        let eig = lobpcg(&a, 1, 300, 1e-7, None).expect("lobpcg");
812        let lambda1 = 2.0 - 2.0 * (PI / ((n + 1) as f64)).cos();
813        assert!(
814            (eig.eigenvalues[0] - lambda1).abs() < 1e-6,
815            "LOBPCG smallest eig {} vs analytic {}",
816            eig.eigenvalues[0],
817            lambda1
818        );
819
820        // Known solution and right-hand side.
821        let x_star: Vec<f64> = (0..n).map(|i| 1.0 + ((i * 7) % 5) as f64).collect();
822        let b = a.matvec(&x_star);
823
824        // (2) Complete IC(k) factor solves the system exactly.
825        let fac = ic_k(&a, n + 1).expect("complete cholesky");
826        let x_ic = fac.apply(&b);
827        for i in 0..n {
828            assert!(
829                (x_ic[i] - x_star[i]).abs() < 1e-7,
830                "IC complete solve mismatch at {i}: {} vs {}",
831                x_ic[i],
832                x_star[i]
833            );
834        }
835
836        // (3) AMG solves the same system to tolerance.
837        let opts = AmgOptions::default();
838        let (x_amg, iters, rel) = amg_solve(&a, &b, &opts, 50, 1e-9).expect("amg solve");
839        assert!(rel < 1e-9, "AMG did not converge: rel = {rel}");
840        assert!(iters < 25, "AMG took too many cycles: {iters}");
841
842        // (4) The two independent solvers agree.
843        for i in 0..n {
844            assert!(
845                (x_amg[i] - x_ic[i]).abs() < 1e-6,
846                "AMG and IC solutions disagree at {i}: {} vs {}",
847                x_amg[i],
848                x_ic[i]
849            );
850        }
851    }
852}