Vanka smoothers

src/GridapSolvers.jl

@@ -1,57 +1,58 @@
 module GridapSolvers
 
-  include("SolverInterfaces/SolverInterfaces.jl")
-  include("MultilevelTools/MultilevelTools.jl")
-  include("BlockSolvers/BlockSolvers.jl")
-  include("PatchBasedSmoothers/PatchBasedSmoothers.jl")
-  include("LinearSolvers/LinearSolvers.jl")
-  include("NonlinearSolvers/NonlinearSolvers.jl")
-
-  using GridapSolvers.SolverInterfaces
-  using GridapSolvers.MultilevelTools
-  using GridapSolvers.BlockSolvers
-  using GridapSolvers.LinearSolvers
-  using GridapSolvers.PatchBasedSmoothers
-  using GridapSolvers.NonlinearSolvers
-
-  # MultilevelTools
-  export get_parts, generate_level_parts, generate_subparts
-
-  export ModelHierarchy, CartesianModelHierarchy, P4estCartesianModelHierarchy
-  export num_levels, get_level, get_level_parts
-  export get_model, get_model_before_redist
-
-  export FESpaceHierarchy
-  export get_fe_space, get_fe_space_before_redist
-  export compute_hierarchy_matrices
-
-  export DistributedGridTransferOperator
-  export RestrictionOperator, ProlongationOperator
-  export setup_transfer_operators
-
-  # BlockSolvers
-  export BlockDiagonalSolver
-
-  # LinearSolvers
-  export JacobiLinearSolver
-  export RichardsonSmoother
-  export SymGaussSeidelSmoother
-  export GMGLinearSolver
-  export BlockDiagonalSmoother
-
-  export ConjugateGradientSolver
-  export IS_GMRESSolver
-  export IS_MINRESSolver
-  export IS_SSORSolver
-
-  export CGSolver
-  export MINRESSolver
-  export GMRESSolver
-  export FGMRESSolver
-
-  # PatchBasedSmoothers
-  export PatchDecomposition
-  export PatchFESpace
-  export PatchBasedLinearSolver
+include("SolverInterfaces/SolverInterfaces.jl")
+include("MultilevelTools/MultilevelTools.jl")
+include("BlockSolvers/BlockSolvers.jl")
+include("PatchBasedSmoothers/PatchBasedSmoothers.jl")
+include("LinearSolvers/LinearSolvers.jl")
+include("NonlinearSolvers/NonlinearSolvers.jl")
+
+using GridapSolvers.SolverInterfaces
+using GridapSolvers.MultilevelTools
+using GridapSolvers.BlockSolvers
+using GridapSolvers.LinearSolvers
+using GridapSolvers.PatchBasedSmoothers
+using GridapSolvers.NonlinearSolvers
+
+# MultilevelTools
+export get_parts, generate_level_parts, generate_subparts
+
+export ModelHierarchy, CartesianModelHierarchy, P4estCartesianModelHierarchy
+export num_levels, get_level, get_level_parts
+export get_model, get_model_before_redist
+
+export FESpaceHierarchy
+export get_fe_space, get_fe_space_before_redist
+export compute_hierarchy_matrices
+
+export DistributedGridTransferOperator
+export RestrictionOperator, ProlongationOperator
+export setup_transfer_operators
+
+# BlockSolvers
+export BlockDiagonalSolver
+
+# LinearSolvers
+export JacobiLinearSolver
+export RichardsonSmoother
+export SymGaussSeidelSmoother
+export GMGLinearSolver
+export BlockDiagonalSmoother
+
+export ConjugateGradientSolver
+export IS_GMRESSolver
+export IS_MINRESSolver
+export IS_SSORSolver
+
+export CGSolver
+export MINRESSolver
+export GMRESSolver
+export FGMRESSolver
+
+# PatchBasedSmoothers
+export PatchDecomposition
+export PatchFESpace
+export PatchBasedLinearSolver
+export Closure
 
 end

src/LinearSolvers/CallbackSolver.jl

@@ -0,0 +1,52 @@
+
+struct CallbackSolver{A,B} <: Algebra.LinearSolver
+  solver :: A
+  callback :: B
+
+  function CallbackSolver(solver::LinearSolver,callback::Function)
+    A = typeof(solver)
+    B = typeof(callback)
+    new{A,B}(solver,callback)
+  end
+end
+
+struct CallbackSolverSS{A,B} <: Algebra.SymbolicSetup
+  solver :: A
+  ss :: B
+end
+
+function Algebra.symbolic_setup(solver::CallbackSolver,mat::AbstractMatrix)
+  ss = Algebra.symbolic_setup(solver.solver,mat)
+  return CallbackSolverSS(solver,ss)
+end
+
+struct CallbackSolverNS{A,B} <: Algebra.NumericalSetup
+  solver :: A
+  ns :: B
+end
+
+function Algebra.numerical_setup(ss::CallbackSolverSS,mat::AbstractMatrix)
+  ns = Algebra.numerical_setup(ss.ss,mat)
+  return CallbackSolverNS(ss.solver,ns)
+end
+
+function Algebra.numerical_setup(ss::CallbackSolverSS,mat::AbstractMatrix,x::AbstractVector)
+  ns = Algebra.numerical_setup(ss.ss,mat,x)
+  return CallbackSolverNS(ss.solver,ns)
+end
+
+function Algebra.numerical_setup!(ns::CallbackSolverNS,mat::AbstractMatrix)
+  Algebra.numerical_setup!(ns.ns,mat)
+  return ns
+end
+
+function Algebra.numerical_setup!(ns::CallbackSolverNS,mat::AbstractMatrix,x::AbstractVector)
+  Algebra.numerical_setup!(ns.ns,mat,x)
+  return ns
+end
+
+function Algebra.solve!(x::AbstractVector,ns::CallbackSolverNS,b::AbstractVector)
+  solve!(x,ns.ns,b)
+  ns.solver.callback(x)
+  return x
+end

src/LinearSolvers/GMGLinearSolvers.jl

@@ -415,7 +415,9 @@ function gmg_work_vectors(smatrices::AbstractVector{<:AbstractMatrix})
   return work_vectors
 end
 
-function apply_GMG_level!(lev::Integer,xh::Union{PVector,Nothing},rh::Union{PVector,Nothing},ns::GMGNumericalSetup)
+function apply_GMG_level!(
+  lev::Integer,xh::Union{<:AbstractVector,Nothing},rh::Union{<:AbstractVector,Nothing},ns::GMGNumericalSetup
+)
   mh = ns.solver.mh
   parts = get_level_parts(mh,lev)
   if i_am_in(parts)

src/LinearSolvers/LinearSolverFromSmoothers.jl

@@ -0,0 +1,50 @@
+
+struct LinearSolverFromSmoother{A} <: Algebra.LinearSolver
+  smoother :: A
+end
+
+struct LinearSolverFromSmootherSS{A,B} <: Algebra.SymbolicSetup
+  smoother :: A
+  smoother_ss :: B
+end
+
+struct LinearSolverFromSmootherNS{A,B,C} <: Algebra.NumericalSetup
+  smoother    :: A
+  smoother_ns :: B
+  caches      :: C
+end
+
+function Gridap.Algebra.symbolic_setup(solver::LinearSolverFromSmoother, mat::AbstractMatrix)
+  ss = symbolic_setup(solver.smoother,mat)
+  return LinearSolverFromSmootherSS(solver.smoother,ss)
+end
+
+function Gridap.Algebra.numerical_setup(ss::LinearSolverFromSmootherSS, mat::AbstractMatrix)
+  ns = numerical_setup(ss.smoother_ss, mat)
+  caches = allocate_in_domain(mat)
+  return LinearSolverFromSmootherNS(ss.smoother,ns,caches)
+end
+
+function Gridap.Algebra.numerical_setup(ss::LinearSolverFromSmootherSS, mat::AbstractMatrix, vec::AbstractVector)
+  ns = numerical_setup(ss.smoother_ss, mat, vec)
+  caches = allocate_in_domain(mat)
+  return LinearSolverFromSmootherNS(ss.smoother,ns,caches)
+end
+
+function Gridap.Algebra.numerical_setup!(ns::LinearSolverFromSmootherNS, mat::AbstractMatrix)
+  numerical_setup!(ns.smoother_ns, mat)
+  return ns
+end
+
+function Gridap.Algebra.numerical_setup!(ns::LinearSolverFromSmootherNS, mat::AbstractMatrix, vec::AbstractVector)
+  numerical_setup!(ns.smoother_ns, mat, vec)
+  return ns
+end
+
+function Gridap.Algebra.solve!(x::AbstractVector,ns::LinearSolverFromSmootherNS, b::AbstractVector)
+  r = ns.caches
+  fill!(x,zero(eltype(x)))
+  copy!(r,b)
+  solve!(x,ns.smoother_ns,r)
+  return x
+end

src/LinearSolvers/LinearSolvers.jl

@@ -18,12 +18,16 @@ using GridapSolvers.MultilevelTools
 using GridapSolvers.SolverInterfaces
 using GridapSolvers.PatchBasedSmoothers
 
+export LinearSolverFromSmoother
 export JacobiLinearSolver
 export RichardsonSmoother
 export SymGaussSeidelSmoother
 export GMGLinearSolver
 export BlockDiagonalSmoother
 export SchurComplementSolver
+export SchwarzLinearSolver
+
+export CallbackSolver
 
 # Wrappers for IterativeSolvers.jl
 export IS_ConjugateGradientSolver
@@ -49,12 +53,18 @@ include("PETSc/ElasticitySolvers.jl")
 include("PETSc/HipmairXuSolvers.jl")
 
 include("IdentityLinearSolvers.jl")
+
+include("LinearSolverFromSmoothers.jl")
 include("JacobiLinearSolvers.jl")
 include("RichardsonSmoothers.jl")
 include("SymGaussSeidelSmoothers.jl")
+
 include("GMGLinearSolvers.jl")
 include("IterativeLinearSolvers.jl")
 include("SchurComplementSolvers.jl")
 include("MatrixSolvers.jl")
+include("SchwarzLinearSolvers.jl")
+
+include("CallbackSolver.jl")
 
 end

src/LinearSolvers/SchwarzLinearSolvers.jl

@@ -0,0 +1,49 @@
+
+# TODO: 
+# - Implement the multiplicative case
+# - Add support for weights/averaging when aggregating the additive case?
+
+struct SchwarzLinearSolver{T,S,A} <: Algebra.LinearSolver
+  local_solvers::A
+  function SchwarzLinearSolver(
+    solver::Union{S,AbstractVector{<:S}};
+    type = :additive
+  ) where S <: Algebra.LinearSolver
+    @check type in (:additive,:multiplicative)
+    @notimplementedif type == :multiplicative # TODO
+    A = typeof(solver)
+    new{type,S,A}(solver)
+  end
+end
+
+struct SchwarzSymbolicSetup{T,S,A,B} <: Algebra.SymbolicSetup
+  solver::SchwarzLinearSolver{T,S,A}
+  local_ss::B
+end
+
+function Algebra.symbolic_setup(s::SchwarzLinearSolver,mat::AbstractMatrix)
+  # TODO: This is where we should compute the comm coloring for the multiplicative case
+  expand(s) = map(m -> s,partition(mat))
+  expand(s::AbstractVector) = s
+
+  local_solvers = expand(s.local_solvers)
+  local_ss = map(symbolic_setup,local_solvers,partition(mat))
+  return SchwarzSymbolicSetup(s,local_ss)
+end
+
+struct SchwarzNumericalSetup{T,S,A,B} <: Algebra.NumericalSetup
+  solver::SchwarzLinearSolver{T,S,A}
+  local_ns::B
+end
+
+function Algebra.numerical_setup(ss::SchwarzSymbolicSetup,mat::PSparseMatrix)
+  local_ns = map(numerical_setup,ss.local_ss,partition(mat))
+  return SchwarzNumericalSetup(ss.solver,local_ns)
+end
+
+function Algebra.solve!(x::PVector,ns::SchwarzNumericalSetup{:additive},b::PVector)
+  map(solve!,partition(x),ns.local_ns,partition(b))
+  assemble!(x) |> wait
+  consistent!(x) |> wait
+  return x
+end

src/MultilevelTools/FESpaceHierarchies.jl

@@ -56,6 +56,16 @@ function _cell_conformity(
   return CellConformity(cell_reffe,conformity)
 end
 
+function _cell_conformity(
+  model::DiscreteModel,
+  reffe::ReferenceFE; 
+  conformity=nothing, kwargs...
+) :: CellConformity
+  cell_reffe = CompressedArray([reffe],fill(one(Int8),num_cells(model)))
+  conformity = Conformity(reffe,conformity)
+  return CellConformity(cell_reffe,conformity)
+end
+
 function _cell_conformity(
   model::GridapDistributed.DistributedDiscreteModel,args...;kwargs...
 ) :: AbstractVector{<:CellConformity}
@@ -134,6 +144,31 @@ function Gridap.MultiField.MultiFieldFESpace(spaces::Vector{<:HierarchicalArray}
   end
 end
 
+# Constant FESpaces
+
+function FESpaces.ConstantFESpace(mh::ModelHierarchy)
+  map(mh) do mhl
+    ConstantFESpace(mhl)
+  end
+end
+
+function FESpaces.ConstantFESpace(mh::MultilevelTools.ModelHierarchyLevel)
+  reffe = ReferenceFE(lagrangian,Float64,0)
+  if has_redistribution(mh)
+    cparts, _ = get_old_and_new_parts(mh.red_glue,Val(false))
+    Vh     = i_am_in(cparts) ? ConstantFESpace(get_model_before_redist(mh)) : nothing
+    Vh_red = ConstantFESpace(get_model(mh))
+    cell_conformity = i_am_in(cparts) ? _cell_conformity(get_model_before_redist(mh),reffe) : nothing
+    cell_conformity_red = _cell_conformity(get_model(mh),reffe)
+  else
+    Vh = ConstantFESpace(get_model(mh))
+    Vh_red = nothing
+    cell_conformity = _cell_conformity(get_model(mh),reffe)
+    cell_conformity_red = nothing
+  end
+  return FESpaceHierarchyLevel(mh.level,Vh,Vh_red,cell_conformity,cell_conformity_red,mh)
+end
+
 # Computing system matrices
 
 function compute_hierarchy_matrices(