fix: improve AMG implementation with better neighbor collection and residual handling

Author: DimitriTimoz

benches/benchmark_iteratives.rs

@@ -185,6 +185,7 @@ fn bench_methods(c: &mut Criterion) {
         // Only run AMG if BENCH_METHOD is unset or matches
         #[cfg(feature = "amg")]
         if bench_method.as_deref().is_none_or(|m| m.eq_ignore_ascii_case("amg")) {
+            use nalgebra_sparse_linalg::iteratives::amg;
             group.bench_with_input(BenchmarkId::new("AMG", n), &n, |be, &_n| {
                 be.iter_batched(
                     || {

src/iteratives/amg/graph.rs

@@ -15,15 +15,11 @@ where
                 let values = row.values();
 
                 // Find the maximum absolute off-diagonal value in the current row i
-                let mut max_abs_off_diag_val = T::zero();
-                for (col_idx, val) in col_indices.iter().zip(values.iter()) {
-                    if *col_idx != i { // Exclude diagonal
-                        let abs_val = val.abs();
-                        if abs_val > max_abs_off_diag_val {
-                            max_abs_off_diag_val = abs_val;
-                        }
-                    }
-                }
+                let max_abs_off_diag_val = col_indices.iter()
+                    .zip(values.iter())
+                    .filter(|(col_idx, _)| **col_idx != i)
+                    .map(|(_, val)| val.abs())
+                    .fold(T::zero(), |acc, val| if val > acc { val } else { acc });
 
                 // If there are no off-diagonal elements, no strong connections can be formed.
                 if max_abs_off_diag_val == T::zero() {
@@ -33,17 +29,13 @@ where
 
                 let threshold_val = max_abs_off_diag_val * theta;
 
-                // Pre-allocate vector for strong neighbors (estimate based on typical sparsity)
-                let mut strong_neighbors = Vec::with_capacity(col_indices.len().min(10));
-                
-                // Filter connections based on the threshold
-                for (j_idx, val) in col_indices.iter().zip(values.iter()) {
-                    // Point j is a strong neighbor of i if |a_ij| >= threshold_val
-                    // And j must not be i itself.
-                    if *j_idx != i && val.abs() >= threshold_val {
-                        strong_neighbors.push(*j_idx);
-                    }
-                }
+                // Collect strong neighbors directly using iterator chains
+                let strong_neighbors: Vec<usize> = col_indices.iter()
+                    .zip(values.iter())
+                    .filter_map(|(j_idx, val)| {
+                        (*j_idx != i && val.abs() >= threshold_val).then_some(*j_idx)
+                    })
+                    .collect();
 
                 result.push(strong_neighbors);
             },

src/iteratives/amg/interpolate.rs

@@ -13,59 +13,73 @@ where
     let n = a.nrows();
     let n_coarse = coarse_of.iter().filter(|&&c| c != usize::MAX).count();
 
+    // Better capacity estimation based on connectivity
+    let f_points = marks.iter().filter(|&&m| matches!(m, Mark::F)).count();
+    let estimated_nnz = n_coarse + f_points * 4; // More accurate estimate
+    
     let mut trip = CooMatrix::new(n, n_coarse);
-    // Pre-estimate capacity to reduce reallocations
-    let estimated_nnz = n + (n - n_coarse) * 3; // Rough estimate
     trip.reserve(estimated_nnz);
 
+    // Pre-allocate reusable vectors
+    let mut c_neighbors = Vec::with_capacity(8);
+    let mut weights = Vec::with_capacity(8);
+
     for i in 0..n {
         match marks[i] {
             Mark::C => {
-                let j = coarse_of[i];
-                trip.push(i, j, N::one());
+                trip.push(i, coarse_of[i], N::one());
             }
             Mark::F => {
-                // Find C-point neighbors in the strength graph
-                let c_neighbors: Vec<usize> = s[i]
-                    .iter()
-                    .copied()
-                    .filter(|&nbr| matches!(marks[nbr], Mark::C))
-                    .collect();
+                c_neighbors.clear();
+                // Collect coarse neighbors
+                for &nbr in &s[i] {
+                    if matches!(marks[nbr], Mark::C) {
+                        c_neighbors.push(nbr);
+                    }
+                }
 
+                // If no coarse neighbors, add a zero entry
                 if c_neighbors.is_empty() {
-                    // Fallback: connect to first coarse point if available
                     if n_coarse > 0 {
                         trip.push(i, 0, N::one());
                     }
                     continue;
                 }
 
-                // Compute interpolation weights
-                if let Some(diag_entry) = a.get_entry(i, i) {
-                    let diag = diag_entry.into_value();
-                    let mut weight_sum = N::zero();
-                    
-                    // First pass: compute weights
-                    let mut weights = Vec::with_capacity(c_neighbors.len());
+                // Early exit if no diagonal entry
+                let Some(diag_entry) = a.get_entry(i, i) else {
+                    let equal_weight = N::one() / N::from_usize(c_neighbors.len()).unwrap();
                     for &nbr in &c_neighbors {
-                        if let Some(a_ij) = a.get_entry(i, nbr) {
-                            let w = -a_ij.into_value() / diag;
-                            weights.push((coarse_of[nbr], w));
-                            weight_sum += w;
-                        }
+                        trip.push(i, coarse_of[nbr], equal_weight);
+                    }
+                    continue;
+                };
+
+                let diag = diag_entry.into_value();
+                let mut weight_sum = N::zero();
+                
+                weights.clear();
+                weights.reserve(c_neighbors.len());
+                
+                // Compute weights in single pass
+                for &nbr in &c_neighbors {
+                    if let Some(a_ij) = a.get_entry(i, nbr) {
+                        let w = -a_ij.into_value() / diag;
+                        weights.push((coarse_of[nbr], w));
+                        weight_sum += w;
                     }
-                    
-                    // Normalize weights to sum to 1 for better stability
-                    if weight_sum != N::zero() {
-                        for (col, w) in weights {
-                            trip.push(i, col, w / weight_sum);
-                        }
-                    } else {
-                        // Fallback: equal weights
-                        let equal_weight = N::one() / N::from_usize(c_neighbors.len()).unwrap();
-                        for &nbr in &c_neighbors {
-                            trip.push(i, coarse_of[nbr], equal_weight);
-                        }
+                }
+                
+                // Add entries with normalized weights
+                if weight_sum != N::zero() {
+                    let inv_sum = N::one() / weight_sum;
+                    for (col, w) in weights.drain(..) {
+                        trip.push(i, col, w * inv_sum);
+                    }
+                } else {
+                    let equal_weight = N::one() / N::from_usize(c_neighbors.len()).unwrap();
+                    for &nbr in &c_neighbors {
+                        trip.push(i, coarse_of[nbr], equal_weight);
                     }
                 }
             }

src/iteratives/amg/level.rs

@@ -136,7 +136,7 @@ impl<N: RealField + Copy> Hierarchy<N> {
         // 5. Post-smooth: x_l gets updated by nu_post smoothing steps
         gauss_seidel::solve_with_initial_guess(&lev.a, b, x, nu_post, tol);
         *residual_buffer = &lev.a * &*x - b;
-        println!("Level {} residual norm: {}", l, residual_buffer.norm());
+        //println!("Level {} residual norm: {}", l, residual_buffer.norm());
     }
 }
 

src/iteratives/amg/mod.rs

@@ -31,7 +31,7 @@ where
     use level::*;
 
     // Pre-compute initial residual
-    let mut residual_buffer = DVector::from(&a * &*x - b);
+    let residual_buffer = DVector::from(&a * &*x - b);
     let hierarchy = setup(a, theta, 100);
 
     // Check if we're already converged
@@ -44,11 +44,15 @@ where
     let adaptive_tol = tol.max(initial_residual_norm * T::from_f64(1e-3).unwrap());
 
     for i in 0..max_iter {
+        let mut residual_buffer = DVector::zeros(b.len());
         hierarchy.vcycle(0, b, x, &mut residual_buffer, adaptive_tol, 1, 1);
 
+        // Use the residual buffer that was updated by vcycle
+        let current_residual = residual_buffer;
+        
         // Check convergence every few iterations to reduce overhead
         if i % 5 == 4 || i == max_iter - 1 {
-            let residual_norm = residual_buffer.amax();
+            let residual_norm = current_residual.amax();
             if residual_norm <= tol {
                 return true;
             }