Block-diagonal smoothers

Project.toml

@@ -5,6 +5,7 @@ version = "0.2.0"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e"
 GridapDistributed = "f9701e48-63b3-45aa-9a63-9bc6c271f355"
@@ -15,6 +16,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [compat]
 ArgParse = "1"

src/GridapSolvers.jl

@@ -26,7 +26,16 @@ module GridapSolvers
   # LinearSolvers
   export JacobiLinearSolver
   export RichardsonSmoother
+  export SymGaussSeidelSmoother
   export GMGLinearSolver
+  export BlockDiagonalSmoother
+
+  export ConjugateGradientSolver
+  export IS_GMRESSolver
+  export IS_MINRESSolver
+  export IS_SSORSolver
+
+  export GMRESSolver
 
   # PatchBasedSmoothers
   export PatchDecomposition

src/LinearSolvers/BlockDiagonalSmoothers.jl

@@ -0,0 +1,80 @@
+struct BlockDiagonalSmoother{A,B} <: Gridap.Algebra.LinearSolver
+  blocks  :: A
+  solvers :: B
+  function BlockDiagonalSmoother(blocks ::AbstractArray{<:AbstractMatrix},
+                                 solvers::AbstractArray{<:Gridap.Algebra.LinearSolver})
+    @check length(blocks) == length(solvers)
+    A = typeof(blocks)
+    B = typeof(solvers)
+    return new{A,B}(blocks,solvers)
+  end
+end
+
+# Constructors 
+
+function BlockDiagonalSmoother(block_mat :: AbstractBlockMatrix,
+                               solvers   :: AbstractArray{<:Gridap.Algebra.LinearSolver})
+  mat_blocks = diag(blocks(block_mat))
+  return BlockDiagonalSmoother(mat_blocks,solvers)
+end
+
+function BlockDiagonalSmoother(biforms :: AbstractArray{<:Function},
+                               trials  :: AbstractArray{<:FESpace},
+                               tests   :: AbstractArray{<:FESpace},
+                               solvers :: AbstractArray{<:Gridap.Algebra.LinearSolver})
+  mat_blocks = compute_block_matrices(biforms,trials,tests)
+  return BlockDiagonalSmoother(mat_blocks,solvers)
+end
+
+function BlockDiagonalSmoother(biforms :: AbstractArray{<:Function},
+                               U       :: Union{MultiFieldFESpace,GridapDistributed.DistributedMultiFieldFESpace},
+                               V       :: Union{MultiFieldFESpace,GridapDistributed.DistributedMultiFieldFESpace},
+                               solvers :: AbstractArray{<:Gridap.Algebra.LinearSolver})
+  return BlockDiagonalSmoother(biforms,[U...],[V...],solvers)
+end
+
+function compute_block_matrices(biforms :: AbstractArray{<:Function},
+                                trials  :: AbstractArray{<:FESpace},
+                                tests   :: AbstractArray{<:FESpace})
+  @check length(biforms) == length(tests) == length(trials)
+  mat_blocks = map(assemble_matrix,biforms,tests,trials)
+  return mat_blocks
+end
+
+# Symbolic and numerical setup
+struct BlockDiagonalSmootherSS{A,B} <: Gridap.Algebra.SymbolicSetup
+  solver   :: A
+  block_ss :: B
+end
+
+function Gridap.Algebra.symbolic_setup(solver::BlockDiagonalSmoother,mat::AbstractMatrix)
+  block_ss = map(symbolic_setup,solver.solvers,solver.blocks)
+  return BlockDiagonalSmootherSS(solver,block_ss)
+end
+
+struct BlockDiagonalSmootherNS{A,B} <: Gridap.Algebra.NumericalSetup
+  solver   :: A
+  block_ns :: B
+end
+
+function Gridap.Algebra.numerical_setup(ss::BlockDiagonalSmootherSS,mat::AbstractMatrix)
+  solver   = ss.solver
+  block_ns = map(numerical_setup,ss.block_ss,solver.blocks)
+  return BlockDiagonalSmootherNS(solver,block_ns)
+end
+
+# Solve
+
+function Gridap.Algebra.solve!(x::AbstractBlockVector,ns::BlockDiagonalSmootherNS,b::AbstractBlockVector)
+  @check blocklength(x) == blocklength(b) == length(ns.block_ns)
+  for (iB,bns) in enumerate(ns.block_ns)
+    xi = x[Block(iB)]
+    bi = b[Block(iB)]
+    solve!(xi,bns,bi)
+  end
+  return x
+end
+
+function LinearAlgebra.ldiv!(x,ns::BlockDiagonalSmootherNS,b)
+  solve!(x,ns,b)
+end

src/LinearSolvers/GMGLinearSolvers.jl

@@ -88,7 +88,7 @@ function setup_finest_level_cache(mh::ModelHierarchy,smatrices::Vector{<:Abstrac
   parts = get_level_parts(mh,1)
   if i_am_in(parts)
     Ah = smatrices[1]
-    rh = pfill(0.0, partition(axes(Ah,2)))
+    rh = allocate_col_vector(Ah)
     cache = rh
   end
   return cache
@@ -156,14 +156,14 @@ function setup_coarsest_solver_cache(mh::ModelHierarchy,coarsest_solver::PETScLi
 end
 
 function allocate_level_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix},lev::Integer)
-  dxh   = pfill(0.0,partition(axes(smatrices[lev],2)))
-  Adxh  = pfill(0.0,partition(axes(smatrices[lev],1)))
+  dxh   = allocate_col_vector(smatrices[lev])
+  Adxh  = allocate_row_vector(smatrices[lev])
 
   cparts = get_level_parts(mh,lev+1)
   if i_am_in(cparts)
     AH  = smatrices[lev+1]
-    rH  = pfill(0.0,partition(axes(AH,2)))
-    dxH = pfill(0.0,partition(axes(AH,2)))
+    rH  = allocate_col_vector(AH)
+    dxH = allocate_col_vector(AH)
   else
     rH  = nothing
     dxH = nothing

src/LinearSolvers/GMRESSolvers.jl

@@ -0,0 +1,113 @@
+
+# GMRES Solver
+struct GMRESSolver <: Gridap.Algebra.LinearSolver
+  m  ::Int
+  Pl ::Gridap.Algebra.LinearSolver
+  tol::Float64
+end
+
+struct GMRESSymbolicSetup <: Gridap.Algebra.SymbolicSetup
+  solver
+end
+
+function Gridap.Algebra.symbolic_setup(solver::GMRESSolver, A::AbstractMatrix)
+  return GMRESSymbolicSetup(solver)
+end
+
+mutable struct GMRESNumericalSetup <: Gridap.Algebra.NumericalSetup
+  solver
+  A
+  Pl_ns
+  caches
+end
+
+function get_gmres_caches(m,A)
+  w = allocate_col_vector(A)
+  V = [allocate_col_vector(A) for i in 1:m+1]
+  Z = [allocate_col_vector(A) for i in 1:m]
+
+  H = zeros(m+1,m)  # Hessenberg matrix
+  g = zeros(m+1)    # Residual vector
+  c = zeros(m)      # Gibens rotation cosines
+  s = zeros(m)      # Gibens rotation sines
+  return (w,V,Z,H,g,c,s)
+end
+
+function Gridap.Algebra.numerical_setup(ss::GMRESSymbolicSetup, A::AbstractMatrix)
+  solver = ss.solver
+  Pl_ns  = numerical_setup(symbolic_setup(solver.Pl,A),A)
+  caches = get_gmres_caches(solver.m,A)
+  return GMRESNumericalSetup(solver,A,Pl_ns,caches)
+end
+
+function Gridap.Algebra.numerical_setup!(ns::GMRESNumericalSetup, A::AbstractMatrix)
+  numerical_setup!(ns.Pl_ns,A)
+  ns.A = A
+end
+
+function Gridap.Algebra.solve!(x::AbstractVector,ns::GMRESNumericalSetup,b::AbstractVector)
+  solver, A, Pl, caches = ns.solver, ns.A, ns.Pl_ns, ns.caches
+  m, tol = solver.m, solver.tol
+  w, V, Z, H, g, c, s = caches
+  println(" > Starting GMRES solver: ")
+
+  # Initial residual
+  mul!(w,A,x); w .= b .- w
+
+  β    = norm(w)
+  iter = 0
+  while (β > tol)
+    println("   > Iteration ", iter," - Residual: ", β)
+    fill!(H,0.0)
+
+    # Arnoldi process
+    fill!(g,0.0); g[1] = β
+    V[1] .= w ./ β
+    j = 1
+    while ( j < m+1 && β > tol )
+      println("      > Inner iteration ", j," - Residual: ", β)
+      # Arnoldi orthogonalization by Modified Gram-Schmidt
+      solve!(Z[j],Pl,V[j])
+      mul!(w,A,Z[j])
+      for i in 1:j
+        H[i,j] = dot(w,V[i])
+        w .= w .- H[i,j] .* V[i]
+      end
+      H[j+1,j] = norm(w)
+      V[j+1] = w ./ H[j+1,j]
+
+      # Update QR
+      for i in 1:j-1
+        γ = c[i]*H[i,j] + s[i]*H[i+1,j]
+        H[i+1,j] = -s[i]*H[i,j] + c[i]*H[i+1,j]
+        H[i,j] = γ
+      end
+
+      # New Givens rotation, update QR and residual
+      c[j], s[j], _ = LinearAlgebra.givensAlgorithm(H[j,j],H[j+1,j])
+      H[j,j] = c[j]*H[j,j] + s[j]*H[j+1,j]; H[j+1,j] = 0.0
+      g[j+1] = -s[j]*g[j]; g[j] = c[j]*g[j]
+
+      β  = abs(g[j+1])
+      j += 1
+    end
+    j = j-1
+
+    # Solve least squares problem Hy = g by backward substitution
+    for i in j:-1:1
+      g[i] = (g[i] - dot(H[i,i+1:j],g[i+1:j])) / H[i,i]
+    end
+
+    # Update solution & residual
+    for i in 1:j
+      x .+= g[i] .* Z[i]
+    end
+    mul!(w,A,x); w .= b .- w
+
+    iter += 1
+  end
+  println("   Exiting GMRES solver.")
+  println("   > Num Iter: ", iter," - Final residual: ", β)
+
+  return x
+end

src/LinearSolvers/Helpers.jl

@@ -0,0 +1,29 @@
+
+# Row/Col vector allocations for serial
+function allocate_row_vector(A::AbstractMatrix{T}) where T
+  return zeros(T,size(A,1))
+end
+
+function allocate_col_vector(A::AbstractMatrix{T}) where T
+  return zeros(T,size(A,2))
+end
+
+# Row/Col vector allocations for parallel
+function allocate_row_vector(A::PSparseMatrix)
+  T = eltype(A)
+  return pfill(zero(T),partition(axes(A,1)))
+end
+
+function allocate_col_vector(A::PSparseMatrix)
+  T = eltype(A)
+  return pfill(zero(T),partition(axes(A,2)))
+end
+
+# Row/Col vector allocations for blocks
+function allocate_row_vector(A::AbstractBlockMatrix)
+  return mortar(map(Aii->allocate_row_vector(Aii),blocks(A)[:,1]))
+end
+
+function allocate_col_vector(A::AbstractBlockMatrix)
+  return mortar(map(Aii->allocate_col_vector(Aii),blocks(A)[1,:]))
+end

src/LinearSolvers/IdentityLinearSolvers.jl

@@ -0,0 +1,26 @@
+
+# Identity solver, for testing purposes 
+struct IdentitySolver <: Gridap.Algebra.LinearSolver
+end
+
+struct IdentitySymbolicSetup <: Gridap.Algebra.SymbolicSetup
+  solver
+end
+
+function Gridap.Algebra.symbolic_setup(s::IdentitySolver,A::AbstractMatrix) 
+  IdentitySymbolicSetup(s)
+end
+
+struct IdentityNumericalSetup <: Gridap.Algebra.NumericalSetup
+  solver
+end
+
+function Gridap.Algebra.numerical_setup(ss::IdentitySymbolicSetup,mat::AbstractMatrix)
+  s = ss.solver
+  return IdentityNumericalSetup(s)
+end
+
+function Gridap.Algebra.solve!(x::AbstractVector,ns::IdentityNumericalSetup,y::AbstractVector)
+  copy!(x,y)
+  return x
+end