Add user-defined near null space to aggregation-based AMG

src/aggregation.jl

@@ -1,5 +1,6 @@
-function smoothed_aggregation(A::TA, _B = nothing,
+function smoothed_aggregation(A::TA,
                         ::Type{Val{bs}}=Val{1};
+                        B = nothing,
                         symmetry = HermitianSymmetry(),
                         strength = SymmetricStrength(),
                         aggregate = StandardAggregation(),
@@ -12,9 +13,8 @@ function smoothed_aggregation(A::TA, _B = nothing,
                         diagonal_dominance = false,
                         keep = false,
                         coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
-
     n = size(A, 1)
-    B = isnothing(_B) ? ones(T,n) : copy(_B)
+    B = isnothing(B) ? ones(T,n) : copy(B)
     @assert size(A, 1) == size(B, 1)
 
     #=max_levels, max_coarse, strength =

src/multilevel.jl

@@ -287,47 +287,7 @@ end
 abstract type AMGAlg end
 
 struct RugeStubenAMG  <: AMGAlg end
-
-"""
-    SmoothedAggregationAMG(B=nothing)
-
-Smoothed Aggregation AMG (Algebraic Multigrid) algorithm configuration.
-
-# Arguments
-- `B::Union{AbstractArray, Nothing}`: The **near null space** in SA-AMG represents the low energy that cannot be attenuated by relaxtion, and thus needs to be perserved across the coarse grid. 
-
-    - `B` can be:
-        - a `Vector` (e.g., for scalar PDEs),
-        - a `Matrix` (e.g., for vector PDEs or systems with multiple equations),
-        - or `nothing`.
-
-# Notes
-If `B` is set to `nothing`, it will be internally defaulted to `B = ones(T, n)`, where `T = eltype(A)` and `n = size(A, 1)`.
-
-# Examples
-
-**Poisson equation (scalar PDE):**
-```julia
-n = size(A, 1)
-B = ones(Float64, n)
-amg = SmoothedAggregationAMG(B)
-```
-
-**Linear elasticity equation in 2d (vector PDE):**
-```julia
-n = size(A, 1)          # Ndof = 2 * number of nodes
-B = zeros(Float64, n, 3) 
-B[1:2:end,   :] = [1.0, 0.0, -y]
-B[2:2:end, :] = [0.0, 1.0,  x]
-amg = SmoothedAggregationAMG(B)
-```
-"""
-struct SmoothedAggregationAMG  <: AMGAlg
-    B::Union{<:AbstractArray,Nothing} # `B` can be `Vector`, `Matrix`, or `nothing`
-    function SmoothedAggregationAMG(B::Union{AbstractArray,Nothing}=nothing)
-        new(B)
-    end
-end
+struct SmoothedAggregationAMG  <: AMGAlg end
 
 function solve(A::AbstractMatrix, b::Vector, s::AMGAlg, args...; kwargs...)
     solt = init(s, A, b, args...; kwargs...)
@@ -337,7 +297,7 @@ function init(::RugeStubenAMG, A, b, args...; kwargs...)
     AMGSolver(ruge_stuben(A; kwargs...), b)
 end
 function init(sa::SmoothedAggregationAMG, A, b; kwargs...)
-    AMGSolver(smoothed_aggregation(A,sa.B; kwargs...), b)
+    AMGSolver(smoothed_aggregation(A; kwargs...), b)
 end
 function solve!(solt::AMGSolver, args...; kwargs...) 
     _solve(solt.ml, solt.b, args...; kwargs...)   

src/precs.jl

@@ -14,8 +14,7 @@ function SmoothedAggregationPreconBuilder(; blocksize = 1, kwargs...)
 end
 
 function (b::SmoothedAggregationPreconBuilder)(A::AbstractSparseMatrixCSC, p)
-    B = get(b.kwargs, :B, nothing)  # extract nns from kwargs, default to `nothing`
-    return (aspreconditioner(smoothed_aggregation(SparseMatrixCSC(A), B, Val{b.blocksize}; b.kwargs...)),I)
+    return (aspreconditioner(smoothed_aggregation(SparseMatrixCSC(A), Val{b.blocksize}; b.kwargs...)), I)
 end
 
 

test/nns_test.jl

@@ -10,12 +10,12 @@
 
     #2. pass `B` as vector
     B = ones(T,n)
-    x_vec = solve(A, b, SmoothedAggregationAMG(B), maxiter = 1, abstol = 1e-6)
+    x_vec = solve(A, b, SmoothedAggregationAMG(), maxiter = 1, abstol = 1e-6;B=B)
     @test x_vec ≈ x_nothing
 
     #3. pass `B` as matrix
     B = ones(T,n,1)
-    x_mat = solve(A, b, SmoothedAggregationAMG(B), maxiter = 1, abstol = 1e-6)
+    x_mat = solve(A, b, SmoothedAggregationAMG(), maxiter = 1, abstol = 1e-6;B=B)
     @test x_mat ≈ x_nothing
 end
 
@@ -194,11 +194,8 @@ end
         @load "lin_elastic_2d.jld2" A b B
         A = SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, A.nzval)
 
-        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(B); log=true, reltol=1e-10)
-        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(); log=true, reltol=1e-10)
-
-        ml = smoothed_aggregation(A, B)
-        @show ml
+        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10;B=B)
+        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10)
 
         println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
         println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
@@ -233,8 +230,8 @@ end
         @test u[end-1] ≈ (P * L^3) / (3 * E * I) # vertical disp. at the end of the beam
 
 
-        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(B); log=true, reltol=1e-10)
-        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(); log=true, reltol=1e-10)
+        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10,B=B)
+        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10)
 
         println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
         println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])

test/runtests.jl

@@ -348,7 +348,6 @@ for sz in [ (10,10), (20,20), (50,50)]
     strategy = KrylovJL_CG(precs = RugeStubenPreconBuilder())
     @test solve(prob, strategy, atol=1.0e-14) ≈ u0 rtol = 1.0e-8
 
-
     strategy = KrylovJL_CG(precs = SmoothedAggregationPreconBuilder())
     @test solve(prob, strategy, atol=1.0e-14) ≈ u0 rtol = 1.0e-8