From 6062c1d10fe0c25031b0db2c7fa2c4be9a292b19 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Thu, 12 Dec 2024 00:36:58 +1100 Subject: [PATCH] Added StaggeredNonlinearFEOpearators --- src/BlockSolvers/BlockFEOperators.jl | 6 +- src/BlockSolvers/BlockSolvers.jl | 4 +- src/BlockSolvers/StaggeredFEOperators.jl | 196 +++++++++++------- .../BlockSolvers/StaggeredFEOperatorsTests.jl | 126 ++++++++--- 4 files changed, 227 insertions(+), 105 deletions(-) diff --git a/src/BlockSolvers/BlockFEOperators.jl b/src/BlockSolvers/BlockFEOperators.jl index 8c872baf..ff5c98eb 100644 --- a/src/BlockSolvers/BlockFEOperators.jl +++ b/src/BlockSolvers/BlockFEOperators.jl @@ -92,7 +92,7 @@ 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]) + (length(range) == 1) ? f[range[1]] : MultiFieldFESpace(f.spaces[range]) end return block_spaces end @@ -104,11 +104,11 @@ function BlockArrays.blocks(f::GridapDistributed.DistributedMultiFieldFESpace{<: if (length(range) == 1) space = f[range[1]] else - global_sf_spaces = f[range] + global_sf_spaces = f.field_fe_space[range] local_sf_spaces = 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) + space = DistributedMultiFieldFESpace(global_sf_spaces,local_mf_spaces,gids,vector_type) end space end diff --git a/src/BlockSolvers/BlockSolvers.jl b/src/BlockSolvers/BlockSolvers.jl index 17ced70f..23d29c89 100644 --- a/src/BlockSolvers/BlockSolvers.jl +++ b/src/BlockSolvers/BlockSolvers.jl @@ -14,6 +14,8 @@ using GridapDistributed using GridapSolvers.MultilevelTools using GridapSolvers.SolverInterfaces +using Gridap.MultiField: BlockSparseMatrixAssembler + using GridapDistributed: to_parray_of_arrays using GridapDistributed: DistributedMultiFieldFESpace, DistributedMultiFieldFEFunction @@ -33,6 +35,6 @@ export BlockDiagonalSolver export BlockTriangularSolver export BlockFEOperator -export StaggeredFEOperator, StaggeredAffineFEOperator, StaggeredFESolver +export StaggeredFEOperator, StaggeredAffineFEOperator, StaggeredNonlinearFEOperator, StaggeredFESolver end diff --git a/src/BlockSolvers/StaggeredFEOperators.jl b/src/BlockSolvers/StaggeredFEOperators.jl index 2d3158e7..46765e4b 100644 --- a/src/BlockSolvers/StaggeredFEOperators.jl +++ b/src/BlockSolvers/StaggeredFEOperators.jl @@ -23,6 +23,20 @@ function MultiField.get_block_ranges(::StaggeredFEOperator{NB,SB}) where {NB,SB} MultiField.get_block_ranges(NB,SB,Tuple(1:sum(SB))) end +# TODO: This is type piracy -> move to Gridap +MultiField.num_fields(space::FESpace) = 1 + +# TODO: We could reuse gids in distributed +function combine_fespaces(spaces::Vector{<:FESpace}) + NB = length(spaces) + SB = Tuple(map(num_fields,spaces)) + sf_spaces = vcat(map(split_fespace,spaces)...) + MultiFieldFESpace(sf_spaces; style = BlockMultiFieldStyle(NB,SB)) +end + +split_fespace(space::FESpace) = [space] +split_fespace(space::MultiFieldFESpaceTypes) = [space...] + function get_solution(op::StaggeredFEOperator{NB,SB}, xh::MultiFieldFEFunction, k) where {NB,SB} r = MultiField.get_block_ranges(op)[k] if isone(length(r)) # SingleField @@ -78,7 +92,6 @@ function Algebra.solve!(xh, solver::StaggeredFESolver{NB}, op::StaggeredFEOperat xh_k = get_solution(op,xh,k) op_k = get_operator(op,xhs,k) xh_k, cache_k = solve!(xh_k,solvers[k],op_k,nothing) - copy!(get_free_dof_values(get_solution(op,xh,k)),get_free_dof_values(xh_k)) xhs, caches, operators = (xhs...,xh_k), (caches...,cache_k), (operators...,op_k) end return xh, (caches,operators) @@ -127,12 +140,25 @@ struct StaggeredAffineFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB} trial :: BlockFESpaceTypes{NB,SB} test :: BlockFESpaceTypes{NB,SB} + @doc """ + function StaggeredAffineFEOperator( + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace}, + [assems :: Vector{<:Assembler}] + ) + + Constructor for a `StaggeredAffineFEOperator` operator, taking in each + equation as a pair of bilinear/linear forms and the corresponding trial/test spaces. + The trial/test spaces can be single or multi-field spaces. + """ function StaggeredAffineFEOperator( - biforms::Vector{<:Function}, - liforms::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace}, - assems::Vector{<:Assembler} + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace}, + assems :: Vector{<:Assembler} = map(SparseMatrixAssembler,tests,trials) ) @assert length(biforms) == length(liforms) == length(trials) == length(tests) == length(assems) trial = combine_fespaces(trials) @@ -140,68 +166,53 @@ struct StaggeredAffineFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB} NB, SB = length(trials), Tuple(map(num_fields,trials)) new{NB,SB}(biforms,liforms,trials,tests,assems,trial,test) end -end -""" + @doc """ function StaggeredAffineFEOperator( - biforms::Vector{<:Function}, - liforms::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace}, - [assems::Vector{<:Assembler}] - ) - -Constructor for a `StaggeredAffineFEOperator` operator, taking in each -equation as a pair of bilinear/linear forms and the corresponding trial/test spaces. -The trial/test spaces can be single or multi-field spaces. -""" -function StaggeredAffineFEOperator( - biforms::Vector{<:Function}, - liforms::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace} -) - assems = map(SparseMatrixAssembler,tests,trials) - return StaggeredAffineFEOperator(biforms,liforms,trials,tests,assems) + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trial :: BlockFESpaceTypes{NB,SB,P}, + test :: BlockFESpaceTypes{NB,SB,P}, + [assem :: BlockSparseMatrixAssembler{NB,NV,SB,P}] + ) where {NB,NV,SB,P} + + Constructor for a `StaggeredAffineFEOperator` operator, taking in each + equation as a pair of bilinear/linear forms and the global trial/test spaces. + """ + function StaggeredAffineFEOperator( + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trial :: BlockFESpaceTypes{NB,SB,P}, + test :: BlockFESpaceTypes{NB,SB,P}, + assem :: BlockSparseMatrixAssembler{NB,NV,SB,P} = SparseMatrixAssembler(trial,test) + ) where {NB,NV,SB,P} + @assert length(biforms) == length(liforms) == NB + @assert P == Tuple(1:sum(SB)) "Permutations not supported" + trials = blocks(trial) + tests = blocks(test) + assems = diag(blocks(assem)) + new{NB,SB}(biforms,liforms,trials,tests,assems,trial,test) + end end FESpaces.get_trial(op::StaggeredAffineFEOperator) = op.trial FESpaces.get_test(op::StaggeredAffineFEOperator) = op.test -# Utils - -MultiField.num_fields(space::FESpace) = 1 - -# TODO: We could reuse gids in distributed -function combine_fespaces(spaces::Vector{<:FESpace}) - NB = length(spaces) - SB = Tuple(map(num_fields,spaces)) - - sf_spaces = vcat(map(split_fespace,spaces)...) - MultiFieldFESpace(sf_spaces; style = BlockMultiFieldStyle(NB,SB)) -end - -split_fespace(space::FESpace) = [space] -split_fespace(space::MultiFieldFESpaceTypes) = [space...] - -function get_operator(op::StaggeredFEOperator{NB}, xhs, k) where NB +function get_operator(op::StaggeredAffineFEOperator{NB}, xhs, k) where NB @assert NB >= k a(uk,vk) = op.biforms[k](xhs,uk,vk) l(vk) = op.liforms[k](xhs,vk) - # uhd = zero(op.trials[k]) - # A, b = assemble_matrix_and_vector(a,l,op.assems[k],op.trials[k],op.tests[k],uhd) - # return AffineOperator(A,b) return AffineFEOperator(a,l,op.trials[k],op.tests[k],op.assems[k]) end -function get_operator!(op_k::AffineFEOperator, op::StaggeredFEOperator{NB}, xhs, k) where NB +function get_operator!(op_k::AffineFEOperator, op::StaggeredAffineFEOperator{NB}, xhs, k) where NB @assert NB >= k A, b = get_matrix(op_k), get_vector(op_k) a(uk,vk) = op.biforms[k](xhs,uk,vk) l(vk) = op.liforms[k](xhs,vk) uhd = zero(op.trials[k]) assemble_matrix_and_vector!(a,l,A,b,op.assems[k],op.trials[k],op.tests[k],uhd) - return AffineOperator(A,b) + return op_k end ############################################################################################ @@ -222,20 +233,32 @@ Such a problem is formulated by a set of residual/jacobian pairs: """ struct StaggeredNonlinearFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB} - jacobians :: Vector{<:Function} residuals :: Vector{<:Function} + jacobians :: Vector{<:Function} trials :: Vector{<:FESpace} tests :: Vector{<:FESpace} assems :: Vector{<:Assembler} trial :: BlockFESpaceTypes{NB,SB} test :: BlockFESpaceTypes{NB,SB} + @doc """ + function StaggeredNonlinearFEOperator( + res :: Vector{<:Function}, + jac :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace} + ) + + Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each + equation as a pair of residual/jacobian forms and the corresponding trial/test spaces. + The trial/test spaces can be single or multi-field spaces. + """ function StaggeredNonlinearFEOperator( - res::Vector{<:Function}, - jac::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace}, - assems::Vector{<:Assembler} + res :: Vector{<:Function}, + jac :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace}, + assems :: Vector{<:Assembler} = map(SparseMatrixAssembler,tests,trials) ) @assert length(res) == length(jac) == length(trials) == length(tests) == length(assems) trial = combine_fespaces(trials) @@ -243,31 +266,50 @@ struct StaggeredNonlinearFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB} NB, SB = length(trials), Tuple(map(num_fields,trials)) new{NB,SB}(res,jac,trials,tests,assems,trial,test) end -end -# TODO: Can be compute jacobians from residuals? - -""" + @doc """ function StaggeredNonlinearFEOperator( - res::Vector{<:Function}, - jac::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace} - ) - -Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each -equation as a pair of residual/jacobian forms and the corresponding trial/test spaces. -The trial/test spaces can be single or multi-field spaces. -""" -function StaggeredNonlinearFEOperator( - res::Vector{<:Function}, - jac::Vector{<:Function}, - trials::Vector{<:FESpace}, - tests::Vector{<:FESpace} -) - assems = map(SparseMatrixAssembler,tests,trials) - return StaggeredNonlinearFEOperator(res,jac,trials,tests,assems) + res :: Vector{<:Function}, + jac :: Vector{<:Function}, + trial :: BlockFESpaceTypes{NB,SB,P}, + test :: BlockFESpaceTypes{NB,SB,P}, + [assem :: BlockSparseMatrixAssembler{NB,NV,SB,P}] + ) where {NB,NV,SB,P} + + Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each + equation as a pair of bilinear/linear forms and the global trial/test spaces. + """ + function StaggeredNonlinearFEOperator( + res :: Vector{<:Function}, + jac :: Vector{<:Function}, + trial :: BlockFESpaceTypes{NB,SB,P}, + test :: BlockFESpaceTypes{NB,SB,P}, + assem :: BlockSparseMatrixAssembler{NB,NV,SB,P} = SparseMatrixAssembler(trial,test) + ) where {NB,NV,SB,P} + @assert length(res) == length(jac) == NB + @assert P == Tuple(1:sum(SB)) "Permutations not supported" + trials = blocks(trial) + tests = blocks(test) + assems = diag(blocks(assem)) + new{NB,SB}(res,jac,trials,tests,assems,trial,test) + end end FESpaces.get_trial(op::StaggeredNonlinearFEOperator) = op.trial FESpaces.get_test(op::StaggeredNonlinearFEOperator) = op.test + +function get_operator(op::StaggeredNonlinearFEOperator{NB}, xhs, k) where NB + @assert NB >= k + jac(uk,duk,dvk) = op.jacobians[k](xhs,uk,duk,dvk) + res(uk,vk) = op.residuals[k](xhs,uk,vk) + return FESpaces.FEOperatorFromWeakForm(res,jac,op.trials[k],op.tests[k],op.assems[k]) +end + +function get_operator!( + op_k::FESpaces.FEOperatorFromWeakForm, op::StaggeredNonlinearFEOperator{NB}, xhs, k +) where NB + @assert NB >= k + jac(uk,duk,dvk) = op.jacobians[k](xhs,uk,duk,dvk) + res(uk,vk) = op.residuals[k](xhs,uk,vk) + return FESpaces.FEOperatorFromWeakForm(res,jac,op.trials[k],op.tests[k],op.assems[k]) +end diff --git a/test/BlockSolvers/StaggeredFEOperatorsTests.jl b/test/BlockSolvers/StaggeredFEOperatorsTests.jl index 093d2c1a..c4a8597d 100644 --- a/test/BlockSolvers/StaggeredFEOperatorsTests.jl +++ b/test/BlockSolvers/StaggeredFEOperatorsTests.jl @@ -5,13 +5,83 @@ using Gridap, GridapDistributed, PartitionedArrays using Gridap.MultiField using BlockArrays using GridapSolvers -using GridapSolvers.BlockSolvers, GridapSolvers.LinearSolvers +using GridapSolvers.BlockSolvers, GridapSolvers.LinearSolvers, GridapSolvers.NonlinearSolvers + +function test_solution(xh,sol,X,dΩ) + N = length(sol) + xh_exact = interpolate(sol,X) + for k in 1:N + eh_k = xh[k] - xh_exact[k] + e_k = sqrt(sum(∫(eh_k * eh_k)dΩ)) + println("Error in field $k: $e_k") + @test e_k < 1e-10 + end +end + +function driver_affine(model,verbose) + order = 1 + reffe = ReferenceFE(lagrangian,Float64,order) + V = FESpace(model,reffe;dirichlet_tags="boundary") + + sol = [x -> x[1], x -> x[2], x -> x[1] + x[2], x -> x[1] - x[2]] + U1 = TrialFESpace(V,sol[1]) + U2 = TrialFESpace(V,sol[2]) + U3 = TrialFESpace(V,sol[3]) + U4 = TrialFESpace(V,sol[4]) + + # Define weakforms + Ω = Triangulation(model) + dΩ = Measure(Ω,3*order) + + a1((),u1,v1) = ∫(u1 * v1)dΩ + l1((),v1) = ∫(sol[1] * v1)dΩ + + a2((u1,),(u2,u3),(v2,v3)) = ∫(u1 * u2 * v2)dΩ + ∫(u3 * v3)dΩ + l2((u1,),(v2,v3)) = ∫(sol[2] * u1 * v2)dΩ + ∫(sol[3] * v3)dΩ + + a3((u1,(u2,u3)),u4,v4) = ∫((u1 + u2) * u4 * v4)dΩ + l3((u1,(u2,u3)),v4) = ∫(sol[4] * (u1 + u2) * v4)dΩ + + # Define solver: Each block will be solved with a CG solver + lsolver = CGSolver(JacobiLinearSolver();rtol=1.e-10,verbose=verbose) + solver = StaggeredFESolver(fill(lsolver,3)) + + # Create operator from full spaces + mfs = BlockMultiFieldStyle(3,(1,2,1)) + X = MultiFieldFESpace([U1,U2,U3,U4];style=mfs) + Y = MultiFieldFESpace([V,V,V,V];style=mfs) + op = StaggeredAffineFEOperator([a1,a2,a3],[l1,l2,l3],X,Y) + xh = solve(solver,op) + test_solution(xh,sol,X,dΩ) + + # Create operator from components + UB1, VB1 = U1, V + UB2, VB2 = MultiFieldFESpace([U2,U3]), MultiFieldFESpace([V,V]) + UB3, VB3 = U4, V + op = StaggeredAffineFEOperator([a1,a2,a3],[l1,l2,l3],[UB1,UB2,UB3],[VB1,VB2,VB3]) + + # Solve and keep caches for reuse + xh = zero(X) + xh, cache = solve!(xh,solver,op); + xh, cache = solve!(xh,solver,op,cache); + test_solution(xh,sol,X,dΩ) + + return true +end + + + +############################################################################################ np = (2,2) ranks = DebugArray(LinearIndices((prod(np),))) - +verbose = i_am_main(ranks) model = CartesianDiscreteModel(ranks,np,(0,1,0,1),(4,4)) +@testset "StaggeredAffineFEOperators" driver_affine(model,verbose) + +model = CartesianDiscreteModel((0,1,0,1),(4,4)) + order = 1 reffe = ReferenceFE(lagrangian,Float64,order) V = FESpace(model,reffe;dirichlet_tags="boundary") @@ -22,36 +92,44 @@ U2 = TrialFESpace(V,sol[2]) U3 = TrialFESpace(V,sol[3]) U4 = TrialFESpace(V,sol[4]) -mfs = BlockMultiFieldStyle(3,(1,2,1)) -X = MultiFieldFESpace([U1,U2,U3,U4];style=mfs) -Y = MultiFieldFESpace([V,V,V,V];style=mfs) - +# Define weakforms Ω = Triangulation(model) -dΩ = Measure(Ω,3*order) +dΩ = Measure(Ω,4*order) + +F(u::Function) = x -> (u(x) + 1) * u(x) +F(u) = (u + 1) * u +dF(u,du) = 2.0 * u * du + du + +j1((),u1,du1,dv1) = ∫(dF(u1,du1) * dv1)dΩ +r1((),u1,v1) = ∫((F(u1) - F(sol[1])) * v1)dΩ -a1((),u1,v1) = ∫(u1 * v1)dΩ -l1((),v1) = ∫(sol[1] * v1)dΩ +j2((u1,),(u2,u3),(du2,du3),(dv2,dv3)) = ∫(u1 * dF(u2,du2) * dv2)dΩ + ∫(dF(u3,du3) * dv3)dΩ +r2((u1,),(u2,u3),(v2,v3)) = ∫(u1 * (F(u2) - F(sol[2])) * v2)dΩ + ∫((F(u3) - F(sol[3])) * v3)dΩ -a2((u1,),(u2,u3),(v2,v3)) = ∫(u1 * u2 * v2)dΩ + ∫(u3 * v3)dΩ -l2((u1,),(v2,v3)) = ∫(sol[2] * u1 * v2)dΩ + ∫(sol[3] * v3)dΩ +j3((u1,(u2,u3)),u4,du4,dv4) = ∫(u3 * dF(u4,du4) * dv4)dΩ +r3((u1,(u2,u3)),u4,v4) = ∫(u3 * (F(u4) - F(sol[4])) * v4)dΩ -a3((u1,(u2,u3)),u4,v4) = ∫((u1 + u2) * u4 * v4)dΩ -l3((u1,(u2,u3)),v4) = ∫(sol[4] * (u1 + u2) * v4)dΩ +# Define solver: Each block will be solved with a LU solver +lsolver = LUSolver() +nlsolver = NewtonSolver(lsolver;rtol=1.e-10,verbose=verbose) +solver = StaggeredFESolver(fill(nlsolver,3)) +# Create operator from full spaces +mfs = BlockMultiFieldStyle(3,(1,2,1)) +X = MultiFieldFESpace([U1,U2,U3,U4];style=mfs) +Y = MultiFieldFESpace([V,V,V,V];style=mfs) +op = StaggeredNonlinearFEOperator([r1,r2,r3],[j1,j2,j3],X,Y) +xh = solve(solver,op) +test_solution(xh,sol,X,dΩ) + +# Create operator from components UB1, VB1 = U1, V UB2, VB2 = MultiFieldFESpace([U2,U3]), MultiFieldFESpace([V,V]) UB3, VB3 = U4, V -op = StaggeredAffineFEOperator([a1,a2,a3],[l1,l2,l3],[UB1,UB2,UB3],[VB1,VB2,VB3]) - -solver = StaggeredFESolver(fill(CGSolver(JacobiLinearSolver();rtol=1.e-10),3)) -xh = solve(solver,op) +op = StaggeredNonlinearFEOperator([r1,r2,r3],[j1,j2,j3],[UB1,UB2,UB3],[VB1,VB2,VB3]) -xh_exact = interpolate(sol,X) -for k in 1:4 - eh_k = xh[k] - xh_exact[k] - e_k = sqrt(sum(∫(eh_k * eh_k)dΩ)) - println("Error in field $k: $e_k") - @test e_k < 1e-10 -end +xh = zero(X) +xh, cache = solve!(xh,solver,op); +test_solution(xh,sol,X,dΩ) end # module \ No newline at end of file