From 434bbb33994d1bdecd832ad5f708253bc2c271f7 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Thu, 28 Dec 2023 14:55:57 +1100 Subject: [PATCH] Added BlockTriangularSolvers --- src/BlockSolvers/BlockTriangularSolvers.jl | 161 ++++++++++++++++++ .../BlockTriangularSolversTests.jl | 46 +++++ 2 files changed, 207 insertions(+) create mode 100644 test/BlockSolvers/BlockTriangularSolversTests.jl diff --git a/src/BlockSolvers/BlockTriangularSolvers.jl b/src/BlockSolvers/BlockTriangularSolvers.jl index e69de29b..a2afb5fd 100644 --- a/src/BlockSolvers/BlockTriangularSolvers.jl +++ b/src/BlockSolvers/BlockTriangularSolvers.jl @@ -0,0 +1,161 @@ +struct BlockTriangularSolver{T,N,A,B,C} <: Gridap.Algebra.LinearSolver + blocks :: A + solvers :: B + coeffs :: C + function BlockTriangularSolver( + blocks :: AbstractMatrix{<:SolverBlock}, + solvers :: AbstractVector{<:Gridap.Algebra.LinearSolver}, + coeffs = fill(1.0,size(blocks)), + half = :upper + ) + N = length(solvers) + @check size(blocks,1) == size(blocks,2) == N + @check size(coeffs,1) == size(coeffs,2) == N + @check half ∈ (:upper,:lower) + + A = typeof(blocks) + B = typeof(solvers) + C = typeof(coeffs) + return new{Val{half},N,A,B,C}(blocks,solvers,coeffs) + end +end + +function BlockTriangularSolver(solvers::AbstractVector{<:Gridap.Algebra.LinearSolver}; + is_nonlinear::Matrix{Bool}=fill(false,(length(solvers),length(solvers))), + coeffs=fill(1.0,size(is_nonlinear)), + half=:upper) + blocks = map(nl -> nl ? NonlinearSystemBlock() : LinearSystemBlock(),is_nonlinear) + return BlockTriangularSolver(blocks,solvers,coeffs,half) +end + +# Symbolic setup + +struct BlockTriangularSolverSS{A,B,C} <: Gridap.Algebra.SymbolicSetup + solver :: A + block_ss :: B + block_caches :: C +end + +function Gridap.Algebra.symbolic_setup(solver::BlockTriangularSolver,mat::AbstractBlockMatrix) + mat_blocks = blocks(mat) + block_caches = map(instantiate_block_cache,solver.blocks,mat_blocks) + block_ss = map(symbolic_setup,solver.solvers,diag(block_caches)) + return BlockTriangularSolverSS(solver,block_ss,block_caches) +end + +function Gridap.Algebra.symbolic_setup(solver::BlockTriangularSolver{T,N},mat::AbstractBlockMatrix,x::AbstractBlockVector) where {T,N} + mat_blocks = blocks(mat) + vec_blocks = blocks(x) + block_caches = map(CartesianIndices(solver.blocks)) do I + instantiate_block_cache(solver.blocks[I],mat_blocks[I],vec_blocks[I[2]]) + end + block_ss = map(symbolic_setup,solver.solvers,diag(block_caches),vec_blocks) + return BlockTriangularSolverSS(solver,block_ss,block_caches) +end + +# Numerical setup + +struct BlockTriangularSolverNS{T,A,B,C,D} <: Gridap.Algebra.NumericalSetup + solver :: A + block_ns :: B + block_caches :: C + work_caches :: D + function BlockTriangularSolverNS( + solver::BlockTriangularSolver{T}, + block_ns,block_caches,work_caches + ) where T + A = typeof(solver) + B = typeof(block_ns) + C = typeof(block_caches) + D = typeof(work_caches) + return new{T,A,B,C,D}(solver,block_ns,block_caches,work_caches) + end +end + +function Gridap.Algebra.numerical_setup(ss::BlockTriangularSolverSS,mat::AbstractBlockMatrix) + solver = ss.solver + block_ns = map(numerical_setup,ss.block_ss,diag(ss.block_caches)) + work_caches = allocate_in_range(mat) + return BlockTriangularSolverNS(solver,block_ns,ss.block_caches,work_caches) +end + +function Gridap.Algebra.numerical_setup(ss::BlockTriangularSolverSS,mat::AbstractBlockMatrix,x::AbstractBlockVector) + solver = ss.solver + vec_blocks = blocks(x) + block_ns = map(numerical_setup,ss.block_ss,diag(ss.block_caches),vec_blocks) + work_caches = allocate_in_range(mat) + return BlockTriangularSolverNS(solver,block_ns,ss.block_caches,work_caches) +end + +function Gridap.Algebra.numerical_setup!(ns::BlockTriangularSolverNS,mat::AbstractBlockMatrix) + solver = ns.solver + mat_blocks = blocks(mat) + block_caches = map(update_block_cache!,ns.block_caches,solver.blocks,mat_blocks) + map(numerical_setup!,ns.block_ns,diag(block_caches)) + return ns +end + +function Gridap.Algebra.numerical_setup!(ns::BlockTriangularSolverNS,mat::AbstractBlockMatrix,x::AbstractBlockVector) + solver = ns.solver + mat_blocks = blocks(mat) + vec_blocks = blocks(x) + block_caches = map(CartesianIndices(solver.blocks)) do I + update_block_cache!(ns.block_caches[I],mat_blocks[I],vec_blocks[I[2]]) + end + map(numerical_setup!,ns.block_ns,diag(block_caches),vec_blocks) + return ns +end + +function Gridap.Algebra.solve!(x::AbstractBlockVector,ns::BlockTriangularSolverNS{Val{:lower}},b::AbstractBlockVector) + @check blocklength(x) == blocklength(b) == length(ns.block_ns) + NB = length(ns.block_ns) + c, w = ns.solver.coeffs, ns.work_caches + mats = ns.block_caches + for iB in 1:NB + # Add lower off-diagonal contributions + wi = w[Block(iB)] + copy!(wi,b[Block(iB)]) + for jB in 1:iB-1 + cij = c[iB,jB] + if abs(cij) > eps(cij) + xj = x[Block(jB)] + mul!(wi,mats[iB,jB],xj,-cij,1.0) + end + end + + # Solve diagonal block + nsi = ns.block_ns[iB] + xi = x[Block(iB)] + solve!(xi,nsi,wi) + end + return x +end + +function Gridap.Algebra.solve!(x::AbstractBlockVector,ns::BlockTriangularSolverNS{Val{:upper}},b::AbstractBlockVector) + @check blocklength(x) == blocklength(b) == length(ns.block_ns) + NB = length(ns.block_ns) + c, w = ns.solver.coeffs, ns.work_caches + mats = ns.block_caches + for iB in NB:-1:1 + # Add upper off-diagonal contributions + wi = w[Block(iB)] + copy!(wi,b[Block(iB)]) + for jB in iB+1:NB + cij = c[iB,jB] + if abs(cij) > eps(cij) + xj = x[Block(jB)] + mul!(wi,mats[iB,jB],xj,-cij,1.0) + end + end + + # Solve diagonal block + nsi = ns.block_ns[iB] + xi = x[Block(iB)] + solve!(xi,nsi,wi) + end + return x +end + +function LinearAlgebra.ldiv!(x,ns::BlockTriangularSolverNS,b) + solve!(x,ns,b) +end diff --git a/test/BlockSolvers/BlockTriangularSolversTests.jl b/test/BlockSolvers/BlockTriangularSolversTests.jl new file mode 100644 index 00000000..2e84624a --- /dev/null +++ b/test/BlockSolvers/BlockTriangularSolversTests.jl @@ -0,0 +1,46 @@ + +using BlockArrays, LinearAlgebra +using Gridap, Gridap.MultiField, Gridap.Algebra +using PartitionedArrays, GridapDistributed +using GridapSolvers, GridapSolvers.BlockSolvers + +np = (2,2) +ranks = with_debug() do distribute + distribute(LinearIndices((prod(np),))) +end + +model = CartesianDiscreteModel(ranks,np,(0,1,0,1),(8,8)) + +reffe = ReferenceFE(lagrangian,Float64,1) +V = FESpace(model,reffe) + +mfs = BlockMultiFieldStyle() +Y = MultiFieldFESpace([V,V];style=mfs) + +Ω = Triangulation(model) +dΩ = Measure(Ω,4) + +sol(x) = sum(x) +a((u1,u2),(v1,v2)) = ∫(u1⋅v1 + u2⋅v2 + u1⋅v2 - u2⋅v1)*dΩ +l((v1,v2)) = ∫(sol⋅v1 - sol⋅v2)*dΩ + +op = AffineFEOperator(a,l,Y,Y) +A, b = get_matrix(op), get_vector(op); + +# Upper +s1 = BlockTriangularSolver([LUSolver(),LUSolver()];half=:upper) +ss1 = symbolic_setup(s1,A) +ns1 = numerical_setup(ss1,A) +numerical_setup!(ns1,A) + +x1 = allocate_in_domain(A); fill!(x1,0.0) +solve!(x1,ns1,b) + +# Lower +s2 = BlockTriangularSolver([LUSolver(),LUSolver()];half=:lower) +ss2 = symbolic_setup(s2,A) +ns2 = numerical_setup(ss2,A) +numerical_setup!(ns2,A) + +x2 = allocate_in_domain(A); fill!(x2,0.0) +solve!(x2,ns2,b)