1use crate::error::{SparseError, SparseResult};
33use crate::host_csr::{HostCsr, dense_solve};
34
35pub struct AmgLevel {
37 pub a: HostCsr,
39 pub p: HostCsr,
42 pub r: HostCsr,
44 pub diag: Vec<f64>,
46}
47
48pub struct AmgHierarchy {
50 pub levels: Vec<AmgLevel>,
53 pub coarse: HostCsr,
55 pub pre_sweeps: usize,
57 pub post_sweeps: usize,
59}
60
61pub struct AmgOptions {
63 pub strength_threshold: f64,
65 pub max_levels: usize,
68 pub max_coarse: usize,
70 pub smooth_omega: f64,
73 pub pre_sweeps: usize,
75 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
92pub 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 if n <= opts.max_coarse || levels.len() + 1 >= opts.max_levels {
127 break;
128 }
129
130 let strength = strength_of_connection(¤t, opts.strength_threshold);
132
133 let (aggregates, num_agg) = aggregate(¤t, &strength);
135 if num_agg == 0 || num_agg >= n {
138 break;
139 }
140
141 let p0 = tentative_prolongator(&aggregates, num_agg, n);
143
144 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(¤t, &diag);
152 let omega = opts.smooth_omega / rho.max(1e-30);
153 let p = smooth_prolongator(¤t, &diag, &p0, omega)?;
154
155 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
179pub fn amg_v_cycle(h: &AmgHierarchy, b: &[f64], x: &mut [f64]) {
186 v_cycle_recursive(h, 0, b, x);
187}
188
189fn v_cycle_recursive(h: &AmgHierarchy, level: usize, b: &[f64], x: &mut [f64]) {
192 if level >= h.levels.len() {
193 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 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 let omega = 2.0 / 3.0;
217
218 jacobi_sweeps(&lvl.a, &lvl.diag, b, x, h.pre_sweeps, omega);
220
221 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 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 v_cycle_recursive(h, level + 1, &r_coarse, &mut e_coarse);
235
236 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 jacobi_sweeps(&lvl.a, &lvl.diag, b, x, h.post_sweeps, omega);
244}
245
246pub 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
295type Strength = Vec<Vec<usize>>;
301
302fn 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
328fn 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 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 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 for i in 0..n {
358 if agg[i] != UNASSIGNED {
359 continue;
360 }
361 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 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
395fn tentative_prolongator(agg: &[usize], num_agg: usize, n: usize) -> HostCsr {
400 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
424fn smooth_prolongator(
429 a: &HostCsr,
430 diag: &[f64],
431 p0: &HostCsr,
432 omega: f64,
433) -> SparseResult<HostCsr> {
434 let ap0 = a.matmul(p0)?;
436 let n = a.nrows;
437 let ncols = p0.ncols;
438
439 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 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 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
498fn spectral_radius_estimate(a: &HostCsr, diag: &[f64]) -> f64 {
501 let n = a.nrows;
502 if n == 0 {
503 return 1.0;
504 }
505 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 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 lambda = dot(&v, &w);
523 for i in 0..n {
524 v[i] = w[i] / nrm;
525 }
526 }
527 lambda.abs().max(1.0)
529}
530
531fn 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
548fn dot(a: &[f64], b: &[f64]) -> f64 {
550 a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum()
551}
552
553fn norm(a: &[f64]) -> f64 {
555 dot(a, a).sqrt()
556}
557
558fn 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 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 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 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 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 for lvl in &h.levels {
660 let ac = if lvl.a.nrows == a.nrows {
661 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 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 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 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 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 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 let a = laplacian_1d(10);
761 let high = strength_of_connection(&a, 0.9);
762 assert!(high.iter().all(|nbrs| nbrs.is_empty()));
764 let low = strength_of_connection(&a, 0.4);
765 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 #[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 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 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 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 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 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}