Merge pull request #103 from termi-official/do/linsolvejl

LinearSolve.jl coarse solve integration

Author: ranjanan

Project.toml

@@ -5,6 +5,7 @@ version = "0.5.1"
 [deps]
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -19,8 +20,9 @@ DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DelimitedFiles", "FileIO", "IterativeSolvers", "JLD2", "Random", "Test"]
+test = ["DelimitedFiles", "FileIO", "IterativeSolvers", "JLD2", "LinearSolve", "Random", "Test"]

src/AlgebraicMultigrid.jl

@@ -2,18 +2,13 @@ module AlgebraicMultigrid
 
 using Reexport
 using LinearAlgebra
+using LinearSolve
 using SparseArrays, Printf
-using Base.Threads
 @reexport import CommonSolve: solve, solve!, init
 using Reexport
 
 using LinearAlgebra: rmul!
 
-# const mul! = A_mul_B!
-
-const MT = false
-const AMG = AlgebraicMultigrid
-
 include("utils.jl")
 export approximate_spectral_radius
 

src/multilevel.jl

@@ -52,6 +52,12 @@ function coarse_b!(m::MultiLevelWorkspace{TX, bs}, n) where {TX, bs}
 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)))
@@ -61,6 +67,43 @@ 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)

src/utils.jl

@@ -99,7 +99,6 @@ end
 find_breakdown(::Type{Float64}) = eps(Float64) * 10^6
 find_breakdown(::Type{Float32}) = eps(Float64) * 10^3
 
-using Base.Threads
 #=function mul!(α::Number, A::SparseMatrixCSC, B::StridedVecOrMat, β::Number, C::StridedVecOrMat)
     A.n == size(B, 1) || throw(DimensionMismatch())
     A.m == size(C, 1) || throw(DimensionMismatch())

test/runtests.jl

@@ -1,6 +1,6 @@
 using SparseArrays, DelimitedFiles, Random
 using Test, LinearAlgebra
-using IterativeSolvers, AlgebraicMultigrid
+using IterativeSolvers, LinearSolve, AlgebraicMultigrid
 import AlgebraicMultigrid: Pinv, Classical
 using JLD2
 using FileIO
@@ -128,6 +128,9 @@ ml = ruge_stuben(A, presmoother = fsmoother,
 x = AlgebraicMultigrid._solve(ml, A * ones(1000))
 @test sum(abs2, x - ones(1000)) < 1e-8
 
+ml = ruge_stuben(A, coarse_solver=AlgebraicMultigrid.LinearSolveWrapper(UMFPACKFactorization()))
+x = AlgebraicMultigrid._solve(ml, A * ones(1000))
+@test sum(abs2, x - ones(1000)) < 1e-7
 
 A = include("randlap.jl")