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

.gitignore

@@ -2,3 +2,4 @@
 *.jl.*.cov
 *.jl.mem
 /Manifest.toml
+.vscode

src/AlgebraicMultigrid.jl

@@ -9,7 +9,7 @@ using Printf
 @reexport import CommonSolve: solve, solve!, init
 using Reexport
 
-using LinearAlgebra: rmul!
+using LinearAlgebra: rmul!, qr
 
 include("utils.jl")
 export approximate_spectral_radius

src/aggregation.jl

@@ -1,5 +1,6 @@
-function smoothed_aggregation(A::TA, 
+function smoothed_aggregation(A::TA,
                         ::Type{Val{bs}}=Val{1};
+                        B = nothing,
                         symmetry = HermitianSymmetry(),
                         strength = SymmetricStrength(),
                         aggregate = StandardAggregation(),
@@ -12,10 +13,9 @@ function smoothed_aggregation(A::TA,
                         diagonal_dominance = false,
                         keep = false,
                         coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
-
     n = size(A, 1)
-    # B = kron(ones(n, 1), eye(1))
-    B = ones(T,n)
+    B = isnothing(B) ? ones(T,n) : copy(B)
+    @assert size(A, 1) == size(B, 1)
 
     #=max_levels, max_coarse, strength =
         levelize_strength_or_aggregation(max_levels, max_coarse, strength)
@@ -70,7 +70,7 @@ function extend_hierarchy!(levels, strength, aggregate, smooth,
     # b = zeros(eltype(A), size(A, 1))
 
     # Improve candidates
-    b = zeros(size(A,1))
+    b = zeros(size(A,1),size(B,2))
     improve_candidates(A, B, b)
     T, B = fit_candidates(AggOp, B)
 
@@ -88,59 +88,39 @@ function extend_hierarchy!(levels, strength, aggregate, smooth,
 end
 construct_R(::HermitianSymmetry, P) = P'
 
-function fit_candidates(AggOp, B, tol = 1e-10)
-
+function fit_candidates(AggOp, B; tol=1e-10)
     A = adjoint(AggOp)
-    n_fine, n_coarse = size(A)
-    n_col = n_coarse
-
-    R = zeros(eltype(B), n_coarse)
-    Qx = zeros(eltype(B), nnz(A))
-    # copy!(Qx, B)
-    for i = 1:size(Qx, 1)
-        Qx[i] = B[i]
-    end
-    # copy!(A.nzval, B)
-    for i = 1:n_col
-        for j in nzrange(A,i)
-            row = A.rowval[j]
-            A.nzval[j] = B[row]
-        end
-    end
-    k = 1
-    for i = 1:n_col
-        norm_i = norm_col(A, Qx, i)
-        threshold_i = tol * norm_i
-        if norm_i > threshold_i
-            scale = 1 / norm_i
-            R[i] = norm_i
-        else
-            scale = 0
-            R[i] = 0
+    n_fine, m = ndims(B) == 1 ? (length(B), 1) : size(B)
+    n_fine2, n_agg = size(A)
+    @assert n_fine2 == n_fine
+    n_coarse = m * n_agg
+    T = eltype(B)
+    Qs = spzeros(T, n_fine, n_coarse)
+    R = zeros(T, n_coarse, m)
+
+    for agg in 1:n_agg
+        rows = A.rowval[A.colptr[agg]:A.colptr[agg+1]-1]
+        M = @view B[rows, :]     # size(rows) × m
+
+
+        F = qr(M)
+        r = min(length(rows), m)
+        Qfull = Matrix(F.Q)
+        Qj = Qfull[:, 1:r]
+        Rj = F.R
+
+        offset = (agg - 1) * m
+
+        for local_i in 1:length(rows), local_j in 1:r
+            val = Qj[local_i, local_j]
+            if abs(val) >= tol
+                Qs[rows[local_i], offset+local_j] = val
+            end
         end
-        for j in nzrange(A, i)
-            row = A.rowval[j]
-            # Qx[row] *= scale
-            #@show k
-            # Qx[k] *= scale
-            # k += 1
-            A.nzval[j] *= scale
-        end
-    end
+        dropzeros!(Qs)
 
-    # SparseMatrixCSC(size(A)..., A.colptr, A.rowval, Qx), R
-    A, R
-end
-function norm_col(A, Qx, i)
-    s = zero(eltype(A))
-    for j in nzrange(A, i)
-        if A.rowval[j] > length(Qx)
-            val = 1
-        else
-            val = Qx[A.rowval[j]]
-        end
-        # val = A.nzval[A.rowval[j]]
-        s += val*val
+        R[offset+1:offset+r, :] .= Rj[1:r, :]
     end
-    sqrt(s)
+
+    return Qs, R
 end

src/classical.jl

@@ -9,6 +9,10 @@ function ruge_stuben(_A::Union{TA, Symmetric{Ti, TA}, Hermitian{Ti, TA}},
                 max_coarse = 10,
                 coarse_solver = Pinv, kwargs...) where {Ti,Tv,bs,TA<:SparseMatrixCSC{Ti,Tv}}
 
+
+    # fails if near null space `B` is provided
+    haskey(kwargs, :B) && kwargs[:B] !== nothing && error("near null space `B` is only supported for smoothed aggregation AMG, not Ruge-Stüben AMG.")
+
     if _A isa Symmetric && Ti <: Real || _A isa Hermitian
         A = _A.data
         symmetric = true

src/multilevel.jl

@@ -187,7 +187,13 @@ Keyword Arguments
 * maxiter::Int64 - maximum number of iterations to execute
 * verbose::Bool - display residual at each iteration
 * log::Bool - return vector of residuals along with solution
+* B::AbstractArray - the **near null space** in SA-AMG, which represents the low energy that cannot be attenuated by relaxtion, and thus needs to be perserved across the coarse grid. 
 
+!!! note
+    `B` can be:
+        - a `Vector` (e.g., for scalar PDEs),
+        - a `Matrix` (e.g., for vector PDEs or systems with multiple equations),
+    If `B` is not provided, it defaults to a vector of ones.
 """
 function _solve(ml::MultiLevel, b::AbstractArray, args...; kwargs...)
     n = length(ml) == 1 ? size(ml.final_A, 1) : size(ml.levels[1].A, 1)
@@ -296,9 +302,9 @@ end
 function init(::RugeStubenAMG, A, b, args...; kwargs...)
     AMGSolver(ruge_stuben(A; kwargs...), b)
 end
-function init(::SmoothedAggregationAMG, A, b; kwargs...)
+function init(sa::SmoothedAggregationAMG, A, b; kwargs...)
     AMGSolver(smoothed_aggregation(A; kwargs...), b)
 end
 function solve!(solt::AMGSolver, args...; kwargs...) 
     _solve(solt.ml, solt.b, args...; kwargs...)   
-end
+end

src/smoother.jl

@@ -29,6 +29,7 @@ end
 function gs!(A, b, x, start, step, stop)
     n = size(A, 1)
     z = zero(eltype(A))
+    @assert size(x,2) == size(b, 2) "x and b must have the same number of columns"
     @inbounds for col in 1:size(x, 2)
         for i in start:step:stop
             rsum = z