Change default solver to SPQR (#131)

Author: termi-official

README.md

@@ -86,6 +86,13 @@ smooth!(x, smoother::S, b)
 Where `S` denotes the smoothers cache which also must hold the matrix A.
 `smooth!` performs relaxation steps updating the current iterate x in-place:  x ← x + S⁻¹(b − A·x).
 
+##  Coarse solver
+
+Users can choose from different coarse solvers.
+* AlgebraicMultigrid.LinearSolveWrapper(alg) where alg is any LinearSolve.jl solver (https://docs.sciml.ai/LinearSolve/stable/basics/algorithm_selection/#Algorithm-Categories)
+* AlgebraicMultigrid.QRSolver
+* AlgebraicMultigrid.Pinv
+
 ## Features and Roadmap
 
 This package currently supports:

src/AlgebraicMultigrid.jl

@@ -28,6 +28,8 @@ include("smoother.jl")
 export GaussSeidel, SymmetricSweep, ForwardSweep, BackwardSweep,
         JacobiProlongation, SOR
 
+include("coarse_solver.jl")
+
 include("multilevel.jl")
 export RugeStubenAMG, SmoothedAggregationAMG
 

src/aggregation.jl

@@ -78,7 +78,7 @@ function smoothed_aggregation(A::TA,
                         diagonal_dominance = false,
                         keep = false,
                         verbose = false,
-                        coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
+                        coarse_solver = _default_coarse_solver(A), kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
 
     @timeit_debug "prologue" begin
 
@@ -94,12 +94,10 @@ function smoothed_aggregation(A::TA,
     end
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
-        prev_n_levels = length(levels)
-        @timeit_debug "extend_hierarchy!" A, B, bsr_flag = extend_hierarchy_sa!(levels, strength, aggregate, smooth,
+        @timeit_debug "extend_hierarchy!" A, B, bsr_flag, stop_coarsening = extend_hierarchy_sa!(levels, strength, aggregate, smooth,
                                 improve_candidates, diagonal_dominance,
                                 keep, A, B, presmoother, postsmoother, symmetry, bsr_flag, verbose)
-        # Break if no coarsening happened (all nodes isolated) or degenerate coarse matrix
-        length(levels) == prev_n_levels && break
+        stop_coarsening && break
         coarse_x!(w, size(A, 1))
         coarse_b!(w, size(A, 1))
         residual!(w, size(A, 1))
@@ -131,7 +129,7 @@ function extend_hierarchy_sa!(levels, strength, aggregate, smooth,
     @timeit_debug "aggregation" AggOp = aggregate(S)
 
     # Cannot coarsen further if no aggregates were formed (e.g. all nodes isolated)
-    size(AggOp, 1) == 0 && return A, B, bsr_flag
+    size(AggOp, 1) == 0 && return A, B, bsr_flag, true
 
     # Improve candidates
     b = zeros(eltype(A), size(A,1),size(B,2))
@@ -140,6 +138,7 @@ function extend_hierarchy_sa!(levels, strength, aggregate, smooth,
 
     @timeit_debug "restriction setup" begin
         P = smooth(A, T, S, B)
+        size(P, 2) == 0 && return A, B, true, true
         R = construct_R(symmetry, P)
     end
 
@@ -154,7 +153,7 @@ function extend_hierarchy_sa!(levels, strength, aggregate, smooth,
     bsr_flag = true
 
     # RAP is the coarse matrix and B is the coarse null space
-    RAP, B, bsr_flag
+    RAP, B, bsr_flag, false
 end
 construct_R(::HermitianSymmetry, P) = P'
 construct_R(::NoSymmetry, P) = P'

src/classical.jl

@@ -12,7 +12,7 @@ function ruge_stuben(A::TA,
                 postsmoother = GaussSeidel(),
                 max_levels = 10,
                 max_coarse = 10,
-                coarse_solver = Pinv, kwargs...) where {Ti,Tv,bs,TA<:SparseMatrixCSC{Ti,Tv}}
+                coarse_solver = _default_coarse_solver(A), 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.")
@@ -22,7 +22,8 @@ function ruge_stuben(A::TA,
     residual!(w, size(A, 1))
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
-        @timeit_debug "extend_hierarchy!" A = extend_hierarchy_rs!(levels, strength, CF, A, presmoother, postsmoother, symmetry)
+        @timeit_debug "extend_hierarchy!" A, stop_coarsening = extend_hierarchy_rs!(levels, strength, CF, A, presmoother, postsmoother, symmetry)
+        stop_coarsening && break
         coarse_x!(w, size(A, 1))
         coarse_b!(w, size(A, 1))
         residual!(w, size(A, 1))
@@ -41,6 +42,7 @@ function extend_hierarchy_rs!(levels, strength, CF, A::SparseMatrixCSC{Ti,Tv}, p
     @timeit_debug "strength" S, T = strength(At)
     @timeit_debug "splitting" splitting = CF(S)
     @timeit_debug "interpolation" P, R = direct_interpolation(At, T, splitting)
+    size(P, 2) == 0 && return A, true
     @timeit_debug "RAP" RAP = R * A * P
 
     @timeit_debug "smoother setup" begin
@@ -49,7 +51,7 @@ function extend_hierarchy_rs!(levels, strength, CF, A::SparseMatrixCSC{Ti,Tv}, p
         push!(levels, Level(A, P, R, pre, post))
     end
 
-    return RAP
+    return RAP, false
 end
 
 function direct_interpolation(At, T, splitting)

src/coarse_solver.jl

@@ -0,0 +1,84 @@
+
+abstract type CoarseSolver end
+
+"""
+    Pinv{T} <: CoarseSolver
+
+Moore-Penrose pseudo inverse coarse solver. Calls `pinv`
+"""
+struct Pinv{T} <: CoarseSolver
+    pinvA::Matrix{T}
+    Pinv{T}(A) where T = new{T}(pinv(Matrix(A)))
+end
+Pinv(A) = Pinv{eltype(A)}(A)
+Base.show(io::IO, p::Pinv) = print(io, "Pinv")
+
+(p::Pinv)(x, b) = mul!(x, p.pinvA, b)
+
+# This one is used internally.
+"""
+    LinearSolveWrapperInternal <: CoarseSolver
+
+Helper to allow the usage of LinearSolve.jl solvers for the coarse-level solve. Constructed via `LinearSolveWrapper`.
+"""
+struct LinearSolveWrapperInternal{LC <: LinearSolve.LinearCache} <: CoarseSolver
+    linsolve::LC
+    function LinearSolveWrapperInternal(A, alg::LinearSolve.SciMLLinearSolveAlgorithm)
+        rhs_tmp = zeros(eltype(A), size(A,1))
+        u_tmp   = zeros(eltype(A), size(A,2))
+        linprob = LinearProblem(A, rhs_tmp; u0 = u_tmp, alias_A = false, alias_b = false)
+        linsolve = init(linprob, alg)
+        new{typeof(linsolve)}(linsolve)
+    end
+end
+
+function (p::LinearSolveWrapperInternal{LC})(x, b) where {LC <: LinearSolve.LinearCache}
+    for i ∈ 1:size(b, 2)
+        # Update right hand side
+        p.linsolve.b = b[:, i]
+        # Solve for x and update
+        x[:, i] = solve!(p.linsolve).u
+    end
+end
+
+function Base.show(io::IO, ml::LinearSolveWrapperInternal)
+    print(io, ml.linsolve.alg)
+end
+
+
+# This one simplifies passing of LinearSolve.jl algorithms into AlgebraicMultigrid.jl as coarse solvers.
+"""
+    LinearSolveWrapper <: CoarseSolver
+
+Helper to allow the usage of LinearSolve.jl solvers for the coarse-level solve.
+"""
+struct LinearSolveWrapper{A <: LinearSolve.SciMLLinearSolveAlgorithm} <: CoarseSolver
+    alg::A
+end
+(p::LinearSolveWrapper)(A::AbstractMatrix) = LinearSolveWrapperInternal(A, p.alg)
+
+
+"""
+    QRSolver{F} <: CoarseSolver
+
+Coarse solver using Julia's built-in factorizations via `qr()`.
+"""
+struct QRSolver{F} <: CoarseSolver
+    factorization::F
+    function QRSolver(A)
+        fact = qr(A)
+        new{typeof(fact)}(fact)
+    end
+end
+Base.show(io::IO, p::QRSolver) = print(io, "QRSolver")
+
+function (solver::QRSolver)(x, b)
+    # Handle multiple RHS efficiently
+    for i ∈ 1:size(b, 2)
+        # Use backslash - Julia's factorizations are optimized for this
+        x[:, i] = solver.factorization \ b[:, i]
+    end
+end
+
+# Guess the best coarse solver based on the matrix type
+_default_coarse_solver(A) = QRSolver

src/multilevel.jl

@@ -58,59 +58,6 @@ function coarse_b!(m::MultiLevelWorkspace{TX, bs}, n) where {TX, bs}
     end
 end
 
-abstract type CoarseSolver end
-
-"""
-    Pinv{T} <: CoarseSolver
-
-Moore-Penrose pseudo inverse coarse solver. Calls `pinv`
-"""
-struct Pinv{T} <: CoarseSolver
-    pinvA::Matrix{T}
-    Pinv{T}(A) where T = new{T}(pinv(Matrix(A)))
-end
-Pinv(A) = Pinv{eltype(A)}(A)
-Base.show(io::IO, p::Pinv) = print(io, "Pinv")
-
-(p::Pinv)(x, b) = mul!(x, p.pinvA, b)
-
-# This one is used internally.
-"""
-    LinearSolveWrapperInternal <: CoarseSolver
-
-Helper to allow the usage of LinearSolve.jl solvers for the coarse-level solve. Constructed via `LinearSolveWrapper`.
-"""
-struct LinearSolveWrapperInternal{LC <: LinearSolve.LinearCache} <: CoarseSolver
-    linsolve::LC
-    function LinearSolveWrapperInternal(A, alg::LinearSolve.SciMLLinearSolveAlgorithm)
-        rhs_tmp = zeros(eltype(A), size(A,1))
-        u_tmp   = zeros(eltype(A), size(A,2))
-        linprob = LinearProblem(A, rhs_tmp; u0 = u_tmp, alias_A = false, alias_b = false)
-        linsolve = init(linprob, alg)
-        new{typeof(linsolve)}(linsolve)
-    end
-end
-
-function (p::LinearSolveWrapperInternal{LC})(x, b) where {LC <: LinearSolve.LinearCache}
-    for i ∈ 1:size(b, 2)
-        # Update right hand side
-        p.linsolve.b = b[:, i]
-        # Solve for x and update
-        x[:, i] = solve!(p.linsolve).u
-    end
-end
-
-# This one simplifies passing of LinearSolve.jl algorithms into AlgebraicMultigrid.jl as coarse solvers.
-"""
-    LinearSolveWrapper <: CoarseSolver
-
-Helper to allow the usage of LinearSolve.jl solvers for the coarse-level solve.
-"""
-struct LinearSolveWrapper <: CoarseSolver
-    alg::LinearSolve.SciMLLinearSolveAlgorithm
-end
-(p::LinearSolveWrapper)(A::AbstractMatrix) = LinearSolveWrapperInternal(A, p.alg)
-
 Base.length(ml::MultiLevel) = length(ml.levels) + 1
 
 function Base.show(io::IO, ml::MultiLevel)

test/nns_test.jl

@@ -218,7 +218,7 @@ end
         x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10,B=B)
         println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
         @test A * x_nns ≈ b
-        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10)
+        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(), coarse_solver=AMG.Pinv, log=true, reltol=1e-10)
         println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
         @test !(A * x_wonns ≈ b) && first(residuals_wonns) > last(residuals_wonns)
 
@@ -249,11 +249,12 @@ end
         u = A \ b
         @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(); log=true, reltol=1e-10, B=B, max_levels=2)
+        # FIXME revisit after https://github.com/JuliaLinearAlgebra/AlgebraicMultigrid.jl/pull/129
+        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(); coarse_solver=AMG.Pinv, log=true, reltol=1e-10, B=B, max_levels=2)
         println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
         @test A * x_nns ≈ b
 
-        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(); log=true, reltol=1e-10, max_levels=2)
+        x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(); coarse_solver=AMG.Pinv, log=true, reltol=1e-10, max_levels=2)
         println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
         @test !(A * x_wonns ≈ b)
 

test/runtests.jl

@@ -140,12 +140,12 @@ x = AlgebraicMultigrid._solve(ml, A * ones(100))
 
 end
 
-@testset "Preconditioning" begin
+@testset "Preconditioning non-SPD problem" begin
 A = include("thing.jl")
 n = size(A, 1)
 smoother = GaussSeidel(ForwardSweep())
 ml = ruge_stuben(A, presmoother = smoother,
-                    postsmoother = smoother)
+                    postsmoother = smoother, coarse_solver = AMG.Pinv)
 p = aspreconditioner(ml)
 b = zeros(n)
 b[1] = 1
@@ -170,7 +170,8 @@ diff = x - [  1.88664780e-16,   2.34982727e-16,   2.33917697e-16,
 @test sum(abs2, diff) < 1e-8
 x = solve(A, b, RugeStubenAMG(); presmoother = smoother, 
                                 postsmoother = smoother,
-                                maxiter = 1, abstol = 1e-12)
+                                maxiter = 1, abstol = 1e-12,
+                                coarse_solver = AMG.Pinv)
 diff = x - [ 0.76347046, -0.5498286 , -0.2705487 , -0.15047352, -0.10248021,
         0.60292674, -0.11497073, -0.08460548, -0.06931461,  0.38230708,
        -0.055664  , -0.04854558, -0.04577031,  0.09964325,  0.01825624,
@@ -197,7 +198,7 @@ diff = x - [ 0.82365077, -0.537589  , -0.30632349, -0.19370186, -0.14773294,
 
 # Symmetric GaussSeidel Smoothing
 
-ml = ruge_stuben(A)
+ml = ruge_stuben(A, coarse_solver = AMG.Pinv)
 p = aspreconditioner(ml)
 
 x = cg(A, b, Pl = p, maxiter = 100_000, reltol = 1e-6)

test/test_regression.jl

@@ -59,8 +59,12 @@ end
 @testset "Regression for issue #56" begin
     # Issue #56
     X = poisson(27_000)+24.0*I
-    ml = ruge_stuben(X)
     b = rand(27_000)
+
+    ml = ruge_stuben(X)
+    @test AlgebraicMultigrid._solve(ml, b, reltol = 1e-10) ≈ X \ b rtol = 1e-10
+
+    ml = smoothed_aggregation(X, strength = SymmetricStrength(0.05))
     @test AlgebraicMultigrid._solve(ml, b, reltol = 1e-10) ≈ X \ b rtol = 1e-10
 end