Merge pull request #49 from gridap/block-solvers

Generic Block solvers

src/BlockSolvers/BlockDiagonalSolvers.jl

@@ -0,0 +1,116 @@
+
+struct BlockDiagonalSolver{N,A,B} <: Gridap.Algebra.LinearSolver
+  blocks  :: A
+  solvers :: B
+  function BlockDiagonalSolver(
+    blocks  :: AbstractVector{<:SolverBlock},
+    solvers :: AbstractVector{<:Gridap.Algebra.LinearSolver}
+    )
+    N = length(solvers)
+    @check length(blocks) == N
+
+    A = typeof(blocks)
+    B = typeof(solvers)
+    return new{N,A,B}(blocks,solvers)
+  end
+end
+
+# Constructors
+
+function BlockDiagonalSolver(solvers::AbstractVector{<:Gridap.Algebra.LinearSolver}; 
+                             is_nonlinear::Vector{Bool}=fill(false,length(solvers)))
+  blocks = map(nl -> nl ? NonlinearSystemBlock() : LinearSystemBlock(),is_nonlinear)
+  return BlockDiagonalSolver(blocks,solvers)
+end
+
+function BlockDiagonalSolver(funcs   :: AbstractArray{<:Function},
+                             trials  :: AbstractArray{<:FESpace},
+                             tests   :: AbstractArray{<:FESpace},
+                             solvers :: AbstractArray{<:Gridap.Algebra.LinearSolver};
+                             is_nonlinear::Vector{Bool}=fill(false,length(solvers)))
+  blocks = map(funcs,trials,tests,is_nonlinear) do f,trial,test,nl
+    nl ? TriformBlock(f,trial,test) : BiformBlock(f,trial,test)
+  end
+  return BlockDiagonalSolver(blocks,solvers)
+end
+
+function BlockDiagonalSolver(mats::AbstractVector{<:AbstractMatrix},
+                             solvers::AbstractVector{<:Gridap.Algebra.LinearSolver})
+  blocks = map(MatrixBlock,mats)
+  return BlockDiagonalSolver(blocks,solvers)
+end
+
+# Symbolic setup
+
+struct BlockDiagonalSolverSS{A,B,C} <: Gridap.Algebra.SymbolicSetup
+  solver       :: A
+  block_ss     :: B
+  block_caches :: C
+end
+
+function Gridap.Algebra.symbolic_setup(solver::BlockDiagonalSolver,mat::AbstractBlockMatrix)
+  mat_blocks   = diag(blocks(mat))
+  block_caches = map(instantiate_block_cache,solver.blocks,mat_blocks)
+  block_ss     = map(symbolic_setup,solver.solvers,block_caches)
+  return BlockDiagonalSolverSS(solver,block_ss,block_caches)
+end
+
+function Gridap.Algebra.symbolic_setup(solver::BlockDiagonalSolver,mat::AbstractBlockMatrix,x::AbstractBlockVector)
+  mat_blocks   = diag(blocks(mat))
+  vec_blocks   = blocks(x)
+  block_caches = map(instantiate_block_cache,solver.blocks,mat_blocks,vec_blocks)
+  block_ss     = map(symbolic_setup,solver.solvers,block_caches,vec_blocks)
+  return BlockDiagonalSolverSS(solver,block_ss,block_caches)
+end
+
+# Numerical setup
+
+struct BlockDiagonalSolverNS{A,B,C} <: Gridap.Algebra.NumericalSetup
+  solver       :: A
+  block_ns     :: B
+  block_caches :: C
+end
+
+function Gridap.Algebra.numerical_setup(ss::BlockDiagonalSolverSS,mat::AbstractBlockMatrix)
+  solver     = ss.solver
+  block_ns   = map(numerical_setup,ss.block_ss,ss.block_caches)
+  return BlockDiagonalSolverNS(solver,block_ns,ss.block_caches)
+end
+
+function Gridap.Algebra.numerical_setup(ss::BlockDiagonalSolverSS,mat::AbstractBlockMatrix,x::AbstractBlockVector)
+  solver     = ss.solver
+  vec_blocks = blocks(x)
+  block_ns   = map(numerical_setup,ss.block_ss,ss.block_caches,vec_blocks)
+  return BlockDiagonalSolverNS(solver,block_ns,ss.block_caches)
+end
+
+function Gridap.Algebra.numerical_setup!(ns::BlockDiagonalSolverNS,mat::AbstractBlockMatrix)
+  solver       = ns.solver
+  mat_blocks   = diag(blocks(mat))
+  block_caches = map(update_block_cache!,ns.block_caches,solver.blocks,mat_blocks)
+  map(numerical_setup!,ns.block_ns,block_caches)
+  return ns
+end
+
+function Gridap.Algebra.numerical_setup!(ns::BlockDiagonalSolverNS,mat::AbstractBlockMatrix,x::AbstractBlockVector)
+  solver       = ns.solver
+  mat_blocks   = diag(blocks(mat))
+  vec_blocks   = blocks(x)
+  block_caches = map(update_block_cache!,ns.block_caches,solver.blocks,mat_blocks,vec_blocks)
+  map(numerical_setup!,ns.block_ns,block_caches,vec_blocks)
+  return ns
+end
+
+function Gridap.Algebra.solve!(x::AbstractBlockVector,ns::BlockDiagonalSolverNS,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::BlockDiagonalSolverNS,b)
+  solve!(x,ns,b)
+end

src/BlockSolvers/BlockFEOperators.jl

@@ -0,0 +1,142 @@
+
+struct BlockFEOperator{NB,SB,P} <: FEOperator
+  global_op    :: FEOperator
+  block_ops    :: Matrix{<:Union{<:FEOperator,Missing,Nothing}}
+  is_nonlinear :: Matrix{Bool}
+end
+
+const BlockFESpaceTypes{NB,SB,P} = Union{<:MultiFieldFESpace{<:BlockMultiFieldStyle{NB,SB,P}},<:GridapDistributed.DistributedMultiFieldFESpace{<:BlockMultiFieldStyle{NB,SB,P}}}
+
+function BlockFEOperator(
+  res::Matrix{<:Union{<:Function,Missing,Nothing}},
+  jac::Matrix{<:Union{<:Function,Missing,Nothing}},
+  trial::BlockFESpaceTypes,
+  test::BlockFESpaceTypes;
+  kwargs...
+)
+  assem = SparseMatrixAssembler(test,trial)
+  return BlockFEOperator(res,jac,trial,test,assem)
+end
+
+function BlockFEOperator(
+  res::Matrix{<:Union{<:Function,Missing,Nothing}},
+  jac::Matrix{<:Union{<:Function,Missing,Nothing}},
+  trial::BlockFESpaceTypes{NB,SB,P},
+  test::BlockFESpaceTypes{NB,SB,P},
+  assem::MultiField.BlockSparseMatrixAssembler{NB,NV,SB,P};
+  is_nonlinear::Matrix{Bool}=fill(true,(NB,NB))
+) where {NB,NV,SB,P}
+  @check size(res,1) == size(jac,1) == NB
+  @check size(res,2) == size(jac,2) == NB
+
+  global_res = residual_from_blocks(NB,SB,P,res)
+  global_jac = jacobian_from_blocks(NB,SB,P,jac)
+  global_op  = FEOperator(global_res,global_jac,trial,test,assem)
+
+  trial_blocks = blocks(trial)
+  test_blocks  = blocks(test)
+  assem_blocks = blocks(assem)
+  block_ops = map(CartesianIndices(res)) do I
+    if !ismissing(res[I]) && !isnothing(res[I])
+      FEOperator(res[I],jac[I],test_blocks[I[1]],trial_blocks[I[2]],assem_blocks[I])
+    end
+  end
+  return BlockFEOperator{NB,SB,P}(global_op,block_ops,is_nonlinear)
+end
+
+# BlockArrays API
+
+BlockArrays.blocks(op::BlockFEOperator) = op.block_ops
+
+# FEOperator API
+
+FESpaces.get_test(op::BlockFEOperator) = get_test(op.global_op)
+FESpaces.get_trial(op::BlockFEOperator) = get_trial(op.global_op)
+Algebra.allocate_residual(op::BlockFEOperator,u) = allocate_residual(op.global_op,u)
+Algebra.residual(op::BlockFEOperator,u) = residual(op.global_op,u)
+Algebra.allocate_jacobian(op::BlockFEOperator,u) = allocate_jacobian(op.global_op,u)
+Algebra.jacobian(op::BlockFEOperator,u) = jacobian(op.global_op,u)
+Algebra.residual!(b::AbstractVector,op::BlockFEOperator,u) = residual!(b,op.global_op,u)
+
+function Algebra.jacobian!(A::AbstractBlockMatrix,op::BlockFEOperator{NB},u) where NB
+  map(blocks(A),blocks(op),op.is_nonlinear) do A,op,nl
+    if nl
+      residual!(A,op,u)
+    end
+  end
+  return A
+end
+
+# Private methods
+
+function residual_from_blocks(NB,SB,P,block_residuals)
+  function res(u,v)
+    block_ranges = MultiField.get_block_ranges(NB,SB,P)
+    block_u = map(r -> (length(r) == 1) ? u[r[1]] : Tuple(u[r]), block_ranges)
+    block_v = map(r -> (length(r) == 1) ? v[r[1]] : Tuple(v[r]), block_ranges)
+    block_contrs = map(CartesianIndices(block_residuals)) do I
+      if !ismissing(block_residuals[I]) && !isnothing(block_residuals[I])
+        block_residuals[I](block_u[I[2]],block_v[I[1]])
+      end
+    end
+    return add_block_contribs(block_contrs)
+  end
+  return res
+end
+
+function jacobian_from_blocks(NB,SB,P,block_jacobians)
+  function jac(u,du,dv)
+    block_ranges = MultiField.get_block_ranges(NB,SB,P)
+    block_u  = map(r -> (length(r) == 1) ?  u[r[1]] : Tuple(u[r]) , block_ranges)
+    block_du = map(r -> (length(r) == 1) ? du[r[1]] : Tuple(du[r]), block_ranges)
+    block_dv = map(r -> (length(r) == 1) ? dv[r[1]] : Tuple(dv[r]), block_ranges)
+    block_contrs = map(CartesianIndices(block_jacobians)) do I
+      if !ismissing(block_jacobians[I]) && !isnothing(block_jacobians[I])
+        block_jacobians[I](block_u[I[2]],block_du[I[2]],block_dv[I[1]])
+      end
+    end
+    return add_block_contribs(block_contrs)
+  end
+  return jac
+end
+
+function add_block_contribs(contrs)
+  c = contrs[1]
+  for ci in contrs[2:end]
+    if !ismissing(ci) && !isnothing(ci)
+      c = c + ci
+    end
+  end
+  return c
+end
+
+function BlockArrays.blocks(a::MultiField.BlockSparseMatrixAssembler)
+  return a.block_assemblers
+end
+
+function BlockArrays.blocks(f::MultiFieldFESpace{<:BlockMultiFieldStyle{NB,SB,P}}) where {NB,SB,P}
+  block_ranges = MultiField.get_block_ranges(NB,SB,P)
+  block_spaces = map(block_ranges) do range
+    (length(range) == 1) ? f[range[1]] : MultiFieldFESpace(f[range])
+  end
+  return block_spaces
+end
+
+function BlockArrays.blocks(f::GridapDistributed.DistributedMultiFieldFESpace{<:BlockMultiFieldStyle{NB,SB,P}}) where {NB,SB,P}
+  block_gids   = blocks(get_free_dof_ids(f))
+
+  block_ranges = MultiField.get_block_ranges(NB,SB,P)
+  block_spaces = map(block_ranges,block_gids) do range, gids
+    if (length(range) == 1) 
+      space = f[range[1]]
+    else
+      global_sf_spaces = f[range]
+      local_sf_spaces  = GridapDistributed.to_parray_of_arrays(map(local_views,global_sf_spaces))
+      local_mf_spaces  = map(MultiFieldFESpace,local_sf_spaces)
+      vector_type = GridapDistributed._find_vector_type(local_mf_spaces,gids)
+      space = MultiFieldFESpace(global_sf_spaces,local_mf_spaces,gids,vector_type)
+    end
+    space
+  end
+  return block_spaces
+end

src/BlockSolvers/BlockSolverInterfaces.jl

@@ -0,0 +1,95 @@
+
+abstract type SolverBlock end
+
+abstract type LinearSolverBlock <: SolverBlock end
+abstract type NonlinearSolverBlock <: SolverBlock end
+
+is_nonlinear(::LinearSolverBlock) = false
+is_nonlinear(::NonlinearSolverBlock) = true
+
+function instantiate_block_cache(block::LinearSolverBlock,mat::AbstractMatrix)
+  @abstractmethod
+end
+function instantiate_block_cache(block::NonlinearSolverBlock,mat::AbstractMatrix,x::AbstractVector)
+  @abstractmethod
+end
+function instantiate_block_cache(block::LinearSolverBlock,mat::AbstractMatrix,x::AbstractVector)
+  instantiate_block_cache(block,mat)
+end
+
+function update_block_cache!(cache,block::LinearSolverBlock,mat::AbstractMatrix)
+  return cache
+end
+function update_block_cache!(cache,block::NonlinearSolverBlock,mat::AbstractMatrix,x::AbstractVector)
+  @abstractmethod
+end
+function update_block_cache!(cache,block::LinearSolverBlock,mat::AbstractMatrix,x::AbstractVector)
+  update_block_cache!(cache,block,mat)
+end
+
+# MatrixBlock
+
+struct MatrixBlock{A} <: LinearSolverBlock
+  mat :: A
+  function MatrixBlock(mat::AbstractMatrix)
+    A = typeof(mat)
+    return new{A}(mat)
+  end
+end
+
+instantiate_block_cache(block::MatrixBlock,::AbstractMatrix) = block.mat
+
+# SystemBlocks
+
+struct LinearSystemBlock <: LinearSolverBlock end
+struct NonlinearSystemBlock <: NonlinearSolverBlock end
+
+instantiate_block_cache(block::LinearSystemBlock,mat::AbstractMatrix) = mat
+instantiate_block_cache(block::NonlinearSystemBlock,mat::AbstractMatrix,::AbstractVector) = mat
+update_block_cache!(cache,block::NonlinearSystemBlock,mat::AbstractMatrix,::AbstractVector) = mat
+
+# BiformBlock/TriformBlock
+struct BiformBlock <: LinearSolverBlock
+  f     :: Function
+  trial :: FESpace
+  test  :: FESpace
+  assem :: Assembler
+  function BiformBlock(f::Function,
+                       trial::FESpace,
+                       test::FESpace,
+                       assem=SparseMatrixAssembler(trial,test))
+    return new(f,trial,test,assem)
+  end
+end
+
+struct TriformBlock <: NonlinearSolverBlock
+  f     :: Function
+  trial :: FESpace
+  test  :: FESpace
+  assem :: Assembler
+  function TriformBlock(f::Function,
+                        trial::FESpace,
+                        test::FESpace,
+                        assem=SparseMatrixAssembler(trial,test))
+    return new(f,trial,test,assem)
+  end
+end
+
+function instantiate_block_cache(block::BiformBlock,mat::AbstractMatrix)
+  return assemble_matrix(block.f,block.assem,block.trial,block.test)
+end
+
+function instantiate_block_cache(block::TriformBlock,mat::AbstractMatrix,x::AbstractVector)
+  uh = FEFunction(block.trial,x)
+  f(u,v) = block.f(uh,u,v)
+  return assemble_matrix(f,block.assem,block.trial,block.test)
+end
+
+function update_block_cache!(cache,block::TriformBlock,mat::AbstractMatrix,x::AbstractVector)
+  uh = FEFunction(block.trial,x)
+  f(u,v) = block.f(uh,u,v)
+  assemble_matrix!(mat,f,block.assem,block.trial,block.test)
+end
+
+# CompositeBlock
+# How do we deal with different sparsity patterns? Not trivial...

src/BlockSolvers/BlockSolvers.jl

@@ -0,0 +1,28 @@
+module BlockSolvers
+  using LinearAlgebra
+  using SparseArrays
+  using SparseMatricesCSR
+  using BlockArrays
+  using IterativeSolvers
+
+  using Gridap
+  using Gridap.Helpers, Gridap.Algebra, Gridap.CellData, Gridap.Arrays, Gridap.FESpaces, Gridap.MultiField
+  using PartitionedArrays
+  using GridapDistributed
+
+  using GridapSolvers.MultilevelTools
+  using GridapSolvers.SolverInterfaces
+
+  include("BlockFEOperators.jl")
+
+  include("BlockSolverInterfaces.jl")
+  include("BlockDiagonalSolvers.jl")
+  include("BlockTriangularSolvers.jl")
+
+  export BlockFEOperator
+
+  export MatrixBlock, LinearSystemBlock, NonlinearSystemBlock, BiformBlock, TriformBlock
+
+  export BlockDiagonalSolver
+  export BlockTriangularSolver
+end