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,4 +1,4 @@
-function smoothed_aggregation(A::TA, 
+function smoothed_aggregation(A::TA, _B = nothing,
                         ::Type{Val{bs}}=Val{1};
                         symmetry = HermitianSymmetry(),
                         strength = SymmetricStrength(),
@@ -14,8 +14,8 @@ function smoothed_aggregation(A::TA,
                         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,1) : 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
-        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
+    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
-    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/multilevel.jl

@@ -287,7 +287,12 @@ end
 abstract type AMGAlg end
 
 struct RugeStubenAMG  <: AMGAlg end
-struct SmoothedAggregationAMG  <: AMGAlg end
+struct SmoothedAggregationAMG  <: AMGAlg
+    B::Union{<:AbstractMatrix,Nothing}
+    function SmoothedAggregationAMG(B::Union{AbstractMatrix,Nothing}=nothing)
+        new(B)
+    end
+end
 
 function solve(A::AbstractMatrix, b::Vector, s::AMGAlg, args...; kwargs...)
     solt = init(s, A, b, args...; kwargs...)
@@ -296,9 +301,9 @@ end
 function init(::RugeStubenAMG, A, b, args...; kwargs...)
     AMGSolver(ruge_stuben(A; kwargs...), b)
 end
-function init(::SmoothedAggregationAMG, A, b; kwargs...)
-    AMGSolver(smoothed_aggregation(A; kwargs...), b)
+function init(sa::SmoothedAggregationAMG, A, b; kwargs...)
+    AMGSolver(smoothed_aggregation(A,sa.B; kwargs...), b)
 end
 function solve!(solt::AMGSolver, args...; kwargs...) 
     _solve(solt.ml, solt.b, args...; kwargs...)   
-end
+end

src/precs.jl

@@ -14,7 +14,8 @@ function SmoothedAggregationPreconBuilder(; blocksize = 1, kwargs...)
 end
 
 function (b::SmoothedAggregationPreconBuilder)(A::AbstractSparseMatrixCSC, p)
-    return (aspreconditioner(smoothed_aggregation(SparseMatrixCSC(A), Val{b.blocksize}; b.kwargs...)), I)
+    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)
 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