diff --git a/scripts/Embedded/MWEs/StaggeredFEStateMaps/2-stagger-state-map.jl b/scripts/Embedded/MWEs/StaggeredFEStateMaps/2-stagger-state-map.jl new file mode 100644 index 0000000..ffe993e --- /dev/null +++ b/scripts/Embedded/MWEs/StaggeredFEStateMaps/2-stagger-state-map.jl @@ -0,0 +1,116 @@ +include("core.jl") +include("extensions.jl") + +model = CartesianDiscreteModel((0,1,0,1),(8,8)) +order = 1 +reffe = ReferenceFE(lagrangian,Float64,order) +Ω = Triangulation(model) + +V_φ = TestFESpace(Ω,reffe) +φh = interpolate(x->x[1]*x[2],V_φ) +V_reg = TestFESpace(Ω,reffe) +U_reg = TrialFESpace(V_reg) + +V = FESpace(Ω,reffe;dirichlet_tags="boundary") + +sol = [x -> x[1], x -> x[1] - x[2]] +U1 = TrialFESpace(V,sol[1]) +U2 = TrialFESpace(V,sol[2]) + +# Define weakforms +dΩ = Measure(Ω,3*order) + +a1((),u1,v1,φ) = ∫(φ * u1 * v1)dΩ +l1((),v1,φ) = ∫(sol[1] * v1)dΩ + +a2((u1,),u2,v2,φ) = ∫(φ * u1 * u2 * v2)dΩ +l2((u1,),v2,φ) = ∫(sol[2] * u1 * v2)dΩ + +# Create operator from components +op = StaggeredAffineFEOperator([a1,a2],[l1,l2],[U1,U2],[V,V]) + +φ_to_u = StaggeredAffineFEStateMap(op,V_φ,U_reg) + +xh, cache = forward_solve!(φ_to_u,φh,nothing) +xh, cache = forward_solve!(φ_to_u,φh,cache) + + + +function _get_solutions(op::StaggeredFEOperator{NB},xh) where NB + map(i->get_solution(op,xh,i),Tuple(1:NB)) +end + +function _get_∂res_k_at_xhi(op::StaggeredAffineFEOperator{NB}, xh_comb, i, k) where NB + @assert NB >= k && 1 <= i < k + ak_at_xhi(xhi,vk) = op.biforms[k]((xh_comb[1:i-1]...,xhi,xh_comb[i+1:end-1]...),xh_comb[k],vk) + lk_at_xhi(xhi,vk) = op.liforms[k]((xh_comb[1:i-1]...,xhi,xh_comb[i+1:end-1]...),vk) + res_k_at_xhi(xhi,vk) = ak_at_xhi(xhi,vk) - lk_at_xhi(xhi,vk) + ∂res_k_at_xhi(vk) = ∇(res_k_at_xhi,[xh_comb[i],vk],1) +end + +F((u1,u2),φ) = ∫(1u1*u1 + 2u2*u2 + 5φ*φ)dΩ +xh_comb = _get_solutions(op,xh) +_dFdxj(j) = ∇((xj->F((xh_comb[1:j-1]...,xj,xh_comb[j+1:end]...),φh)))(xh_comb[j]) + +_biforms, _liforms, _spaces, _assmes, _solvers = φ_to_u.biforms, + φ_to_u.liforms, φ_to_u.spaces, φ_to_u.assems, φ_to_u.solvers +a_at_φ = map(a->((xhs,uk,vk) -> a(xhs,uk,vk,φh)),_biforms) +l_at_φ = map(l->((xhs,vk) -> l(xhs,vk,φh)),_liforms) +_op = StaggeredAffineFEOperator(a_at_φ,l_at_φ,_spaces.trials,_spaces.tests,_assmes.assmes) + +# k = 2 +a_adj2((),λ2,Λ2) = _op.biforms[2](xh_comb[1:end-1],Λ2,λ2) +l_adj2((),Λ2) = _dFdxj(2) +# k = 1 +a_adj1((λ2,),λ1,Λ1) = _op.biforms[1](xh_comb[1:end-2],Λ1,λ1) +∂res_2_at_xh1 = _get_∂res_k_at_xhi(_op,xh_comb,1,2) +l_adj1((λ2,),Λ1) = _dFdxj(1) - ∂res_2_at_xh1(λ2) + +a_adjs = [a_adj2,a_adj1] +l_adjs = [l_adj2,l_adj1] +# op_adjoint = StaggeredAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.tests),reverse(_spaces.trials)) +# op_adjoint = StaggeredAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.trials),reverse(_spaces.tests)) + +op_adjoint = StaggeredAdjointAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.trials),reverse(_spaces.tests)) + +λh = zero(op_adjoint.test) +λh, cache = solve!(λh,φ_to_u.solvers.adjoint_solver,op_adjoint,nothing); +λh, cache = solve!(λh,φ_to_u.solvers.adjoint_solver,op_adjoint,cache); + +# Analytic adjoint +∂F∂u1(du1,(u1,u2)) = ∫(2du1*u1)dΩ +∂F∂u2(du2,(u1,u2)) = ∫(4u2*du2)dΩ + +∂F∂u1_vec = assemble_vector(du->∂F∂u1(du,xh_comb),U1) +∂F∂u1_AD = assemble_vector(_dFdxj(1),U1) +@assert norm(∂F∂u1_vec - ∂F∂u1_AD) < 1.e-15 + +∂F∂u2_vec = assemble_vector(du->∂F∂u2(du,xh_comb),U2) +∂F∂u2_AD = assemble_vector(_dFdxj(2),U2) +@assert norm(∂F∂u2_vec - ∂F∂u2_AD) < 1.e-15 + +norm(∂F∂u2_vec - get_vector(cache[2][1])) + +# Stiffness matrices +∂R1∂u1((),u1,du1,dv1) = a1((),dv1,du1,φh) +∂R1∂u1_mat = assemble_matrix((du,dv)->∂R1∂u1((),xh_comb[1],du,dv),V,U1) + +∂R2∂u2((u1,),u2,du2,dv2) = a2((u1,),dv2,du2,φh) +∂R2∂u2_mat = assemble_matrix((du,dv)->∂R2∂u2((xh_comb[1],),xh_comb[2],du,dv),V,U2) + +# Operator order: λ2 problem, λ1 problem +adjoint_op2 = cache[2][1] +adjoint_op1 = cache[2][2] + +norm(get_matrix(adjoint_op1) - ∂R1∂u1_mat,Inf) +norm(get_matrix(adjoint_op2) - ∂R2∂u2_mat,Inf) + +# RHS & solutions +λ2 = ∂R2∂u2_mat\∂F∂u2_vec +norm(λ2 - get_free_dof_values(get_solution(op,λh,1)),Inf)/norm(λ1,Inf) + +∂R2∂u1(du1,(u1,),u2,v2) = ∫(φh * du1 * u2 * v2)dΩ - ∫(sol[2] * du1 * v2)dΩ +∂R2∂u1_vec = assemble_vector(du->∂R2∂u1(du,(xh_comb[1],),xh_comb[2],FEFunction(U2,λ2)),U1) + +λ1 = ∂R1∂u1_mat\(∂F∂u1_vec - ∂R2∂u1_vec) +norm(λ1 - get_free_dof_values(get_solution(op_adjoint,λh,2)),Inf)/norm(λ1,Inf) \ No newline at end of file diff --git a/scripts/Embedded/MWEs/StaggeredFEStateMaps/3-stagger-state-map.jl b/scripts/Embedded/MWEs/StaggeredFEStateMaps/3-stagger-state-map.jl new file mode 100644 index 0000000..dff6049 --- /dev/null +++ b/scripts/Embedded/MWEs/StaggeredFEStateMaps/3-stagger-state-map.jl @@ -0,0 +1,152 @@ +include("core.jl") +include("extensions.jl") + +model = CartesianDiscreteModel((0,1,0,1),(8,8)) +order = 1 +reffe = ReferenceFE(lagrangian,Float64,order) +Ω = Triangulation(model) + +V_φ = TestFESpace(Ω,reffe) +φh = interpolate(x->(x[1]+1)/2,V_φ) +V_reg = TestFESpace(Ω,reffe) +U_reg = TrialFESpace(V_reg) + +V = FESpace(Ω,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 +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Ω + +# 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]) + +φ_to_u = StaggeredAffineFEStateMap(op,V_φ,U_reg) + +xh, cache = forward_solve!(φ_to_u,φh,nothing) +xh, cache = forward_solve!(φ_to_u,φh,cache) + +# function _fix_ith_residual_at_jth_xhs(op::StaggeredAffineFEOperator{NB}, xhs, φh, i, j) where NB +# @assert NB >= i +# # i == j +# a(uk,vk) = op.biforms[i](xhs,uk,vk,φh) +# l(vk) = op.liforms[i](xhs,vk,φh) +# return (uk,vk) -> a(uk,vk) - l(vk) +# end + +function _get_solutions(op::StaggeredFEOperator{NB},xh) where NB + map(i->get_solution(op,xh,i),Tuple(1:NB)) +end + +function _get_∂res_k_at_xhi(op::StaggeredAffineFEOperator{NB}, xh_comb, i, k) where NB + @assert NB >= k && 1 <= i < k + ak_at_xhi(xhi,vk) = op.biforms[k]((xh_comb[1:i-1]...,xhi,xh_comb[i+1:end-1]...),xh_comb[k],vk) + lk_at_xhi(xhi,vk) = op.liforms[k]((xh_comb[1:i-1]...,xhi,xh_comb[i+1:end-1]...),vk) + res_k_at_xhi(xhi,vk) = ak_at_xhi(xhi,vk) - lk_at_xhi(xhi,vk) + ∂res_k_at_xhi(vk) = ∇(res_k_at_xhi,[xh_comb[i],vk],1) +end + +F((u1,(u2,u3),u4),φ) = ∫(1u1*u1 + 2u2*u2 + 3u3*u3 + 4u4*u4 + 5φ*φ)dΩ +xh_comb = _get_solutions(op,xh) +_dFdxj(j) = ∇((xj->F((xh_comb[1:j-1]...,xj,xh_comb[j+1:end]...),φh)))(xh_comb[j]) + +_biforms, _liforms, _spaces, _assmes, _solvers = φ_to_u.biforms, + φ_to_u.liforms, φ_to_u.spaces, φ_to_u.assems, φ_to_u.solvers +a_at_φ = map(a->((xhs,uk,vk) -> a(xhs,uk,vk,φh)),_biforms) +l_at_φ = map(l->((xhs,vk) -> l(xhs,vk,φh)),_liforms) +_op = StaggeredAffineFEOperator(a_at_φ,l_at_φ,_spaces.trials,_spaces.tests,_assmes.assmes) + +# k = 3 +a_adj3((),λ3,Λ3) = _op.biforms[3](xh_comb[1:end-1],Λ3,λ3) +l_adj3((),Λ3) = _dFdxj(3) +# k = 2 +a_adj2((λ3,),λ2,Λ2) = _op.biforms[2](xh_comb[1:end-2],Λ2,λ2) +∂res_3_at_xh2 = _get_∂res_k_at_xhi(_op,xh_comb,2,3) +l_adj2((λ3,),Λ2) = _dFdxj(2) - ∂res_3_at_xh2(λ3) +# k = 1 +a_adj1((λ3,λ2),λ1,Λ1) = _op.biforms[1](xh_comb[1:end-3],Λ1,λ1) +∂res_3_at_xh1 = _get_∂res_k_at_xhi(_op,xh_comb,1,3) +∂res_2_at_xh1 = _get_∂res_k_at_xhi(_op,xh_comb,1,2) +l_adj1((λ3,λ2),Λ1) = _dFdxj(1) - ∂res_3_at_xh1(λ3) - ∂res_2_at_xh1(λ2) + +a_adjs = [a_adj3,a_adj2,a_adj1] +l_adjs = [l_adj3,l_adj2,l_adj1] +# op_adjoint = StaggeredAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.tests),reverse(_spaces.trials)) +# op_adjoint = StaggeredAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.trials),reverse(_spaces.tests)) +op_adjoint = StaggeredAdjointAffineFEOperator(a_adjs,l_adjs,reverse(_spaces.trials),reverse(_spaces.tests)) + +λh = zero(op_adjoint.test) +λh, cache = solve!(λh,φ_to_u.solvers.adjoint_solver,op_adjoint,nothing); + +# Analytic adjoint +∂F∂u1(du1,(u1,(u2,u3),u4)) = ∫(2du1*u1)dΩ +∂F∂u23((du2,du3),(u1,(u2,u3),u4)) = ∫(4u2*du2 + 6u3*du3)dΩ +∂F∂u4(du4,(u1,(u2,u3),u4)) = ∫(8du4*u4)dΩ + +∂F∂u1_vec = assemble_vector(du->∂F∂u1(du,xh_comb),UB1) +∂F∂u1_AD = assemble_vector(_dFdxj(1),UB1) +@assert norm(∂F∂u1_vec - ∂F∂u1_AD) < 1.e-15 + +∂F∂u23_vec = assemble_vector(du->∂F∂u23(du,xh_comb),UB2) +∂F∂u23_AD = assemble_vector(_dFdxj(2),UB2) +@assert norm(∂F∂u23_vec - ∂F∂u23_AD) < 1.e-15 + +∂F∂u4_vec = assemble_vector(du->∂F∂u4(du,xh_comb),UB3) +∂F∂u4_AD = assemble_vector(_dFdxj(3),UB3) +@assert norm(∂F∂u4_vec - ∂F∂u4_AD) < 1.e-15 + +norm(∂F∂u4_vec - get_vector(cache[2][1])) + +# Stiffness matrices +∂R1∂u1((),u1,du1,dv1) = a1((),dv1,du1,φh) +∂R1∂u1_mat = assemble_matrix((du,dv)->∂R1∂u1((),xh_comb[1],du,dv),VB1,UB1) + +∂R2∂u2((u1,),(u2,u3),(du2,du3),(dv2,dv3)) = a2((u1,),(dv2,dv3),(du2,du3),φh) +∂R2∂u2_mat = assemble_matrix((du,dv)->∂R2∂u2((xh_comb[1],),xh_comb[2],du,dv),VB2,UB2) + +∂R3∂u3((u1,(u2,u3)),u4,du4,dv4) = a3((u1,(u2,u3)),dv4,du4,φh) +∂R3∂u3_mat = assemble_matrix((du,dv)->∂R3∂u3((xh_comb[1:2]...,),xh_comb[3],du,dv),VB3,UB3) + +# Operator order: λ3 problem, λ2 problem, λ1 problem +adjoint_op3 = cache[2][1] +adjoint_op2 = cache[2][2] +adjoint_op1 = cache[2][3] + +norm(get_matrix(adjoint_op1) - ∂R1∂u1_mat,Inf) +norm(get_matrix(adjoint_op2) - ∂R2∂u2_mat,Inf) +norm(get_matrix(adjoint_op3) - ∂R3∂u3_mat,Inf) + +# RHS & solutions +λ3 = ∂R3∂u3_mat\∂F∂u4_vec +norm(λ3 - get_free_dof_values(get_solution(op,λh,1)),Inf)/norm(λ3,Inf) + +∂R3∂u23((du2,du3),(u1,(u2,u3)),u4,v4) = ∫(φh * du2 * u4 * v4)dΩ - ∫(sol[4] * du2 * v4)dΩ +∂R3∂u23_vec = assemble_vector(du->∂R3∂u23(du,(xh_comb[1:2]...,),xh_comb[3],FEFunction(UB3,λ3)),UB2) + +λ2 = ∂R2∂u2_mat\(∂F∂u23_vec - ∂R3∂u23_vec) +get_free_dof_values(get_solution(op_adjoint,λh,2)) +norm(λ2 - get_free_dof_values(get_solution(op_adjoint,λh,2)),Inf)/norm(λ2,Inf) + +∂R3∂u1(du1,(u1,(u2,u3)),u4,v4) = ∫(φh * du1 * u4 * v4)dΩ - ∫(sol[4] * du1 * v4)dΩ +∂R3∂u1_vec = assemble_vector(du->∂R3∂u1(du,(xh_comb[1:2]...,),xh_comb[3],FEFunction(UB3,λ3)),UB1) +∂R2∂u1(du1,(u1,),(u2,u3),(v2,v3)) = ∫(φh * du1 * u2 * v2)dΩ - ∫(sol[2] * du1 * v2)dΩ +∂R2∂u1_vec = assemble_vector(du->∂R2∂u1(du,(xh_comb[1],),xh_comb[2],FEFunction(UB2,λ2)),UB1) + +λ1 = ∂R1∂u1_mat\(∂F∂u1_vec - ∂R2∂u1_vec - ∂R3∂u1_vec) +norm(λ1 - get_free_dof_values(get_solution(op_adjoint,λh,3)),Inf)/norm(λ1,Inf) \ No newline at end of file diff --git a/scripts/Embedded/MWEs/StaggeredFEStateMaps/core.jl b/scripts/Embedded/MWEs/StaggeredFEStateMaps/core.jl new file mode 100644 index 0000000..dda05d5 --- /dev/null +++ b/scripts/Embedded/MWEs/StaggeredFEStateMaps/core.jl @@ -0,0 +1,111 @@ +using GridapTopOpt +using Gridap +using Gridap.Algebra, Gridap.Geometry, Gridap.MultiField, Gridap.Arrays, + Gridap.FESpaces, Gridap.CellData +using GridapTopOpt: AbstractFEStateMap + +using BlockArrays +using GridapSolvers +using GridapSolvers.BlockSolvers, GridapSolvers.LinearSolvers, GridapSolvers.NonlinearSolvers +using GridapSolvers.BlockSolvers: combine_fespaces, get_solution + +import GridapTopOpt: forward_solve!,dRdφ,update_adjoint_caches!,adjoint_solve!,get_state,get_spaces,get_assemblers, + get_trial_space,get_test_space + +struct StaggeredAffineFEStateMap{NB,SB,A,B,C} <: AbstractFEStateMap + biforms :: Vector{<:Function} + liforms :: Vector{<:Function} + spaces :: A + assems :: B + solvers :: C + + function StaggeredAffineFEStateMap( + op :: StaggeredAffineFEOperator{NB,SB}, + V_φ :: FESpace, + U_reg :: FESpace; + assems_adjoint :: Vector{<:Assembler} = op.assems, + assem_deriv :: Assembler = SparseMatrixAssembler(U_reg,U_reg), + solver :: StaggeredFESolver{NB} = StaggeredFESolver(fill(LUSolver(),length(op.biforms))), + adjoint_solver :: StaggeredFESolver{NB} = StaggeredFESolver(fill(LUSolver(),length(op.biforms))) + ) where {NB,SB} + + spaces = (;trials=op.trials,tests=op.tests,aux_space=V_φ,deriv_space=U_reg, + trial=op.trial, test=op.test) + assems = (;assmes=op.assems,assem_deriv,assems_adjoint) + _solvers = (;solver,adjoint_solver) + A,B,C = typeof(spaces), typeof(assems), typeof(_solvers) + new{NB,SB,A,B,C}(op.biforms,op.liforms,spaces,assems,_solvers) + end +end + +# get_state(m::StaggeredAffineFEStateMap) = isdefined(m,:fwd_caches) ? FEFunction(m.fwd_caches.X,m.fwd_caches.x) : +# error("fwd_caches not defined, call `forward_solve!` to create the cache") +get_state(::StaggeredAffineFEStateMap) = error("This method has been deprecated") +get_spaces(m::StaggeredAffineFEStateMap) = m.spaces +get_assemblers(m::StaggeredAffineFEStateMap) = m.assems + +function forward_solve!(φ_to_u::StaggeredAffineFEStateMap,φh,::Nothing) + biforms, liforms, spaces, assmes, solvers = φ_to_u.biforms, + φ_to_u.liforms, φ_to_u.spaces, φ_to_u.assems, φ_to_u.solvers + a_at_φ = map(a->((xhs,uk,vk) -> a(xhs,uk,vk,φh)),biforms) + l_at_φ = map(l->((xhs,vk) -> l(xhs,vk,φh)),liforms) + op = StaggeredAffineFEOperator(a_at_φ,l_at_φ,spaces.trials,spaces.tests,assmes.assmes) + X = combine_fespaces(spaces.trials) + + xh = zero(X) + xh, cache = solve!(xh,solvers.solver,op); + + x = get_free_dof_values(xh) + return xh, (x,X,op,cache) +end + +function forward_solve!(φ_to_u::StaggeredAffineFEStateMap,φh,cache) + biforms, liforms, spaces, assmes, solvers = φ_to_u.biforms, + φ_to_u.liforms, φ_to_u.spaces, φ_to_u.assems, φ_to_u.solvers + a_at_φ = map(a->((xhs,uk,vk) -> a(xhs,uk,vk,φh)),biforms) + l_at_φ = map(l->((xhs,vk) -> l(xhs,vk,φh)),liforms) + op = StaggeredAffineFEOperator(a_at_φ,l_at_φ,spaces.trials,spaces.tests,assmes.assmes) + + x, X, _, last_cache = cache + xh, cache_updated = solve!(FEFunction(X,x),solvers.solver,op,last_cache); + + return xh, (x,X,op,cache_updated) +end + +function dRdφ(φ_to_u::StaggeredAffineFEStateMap{NB},uhs,λhs,φh) where NB + biforms, liforms, spaces, assmes = φ_to_u.biforms, φ_to_u.liforms, φ_to_u.spaces, φ_to_u.assems + dummy_op = StaggeredAffineFEOperator(biforms,liforms,spaces.trials,spaces.tests,assmes.assmes) + xhs, ∂Rs∂φ = (), () + for k in 1:NB + xh_k = get_solution(op,uhs,k) + λh_k = get_solution(op,λhs,k) + _a(uk,vk,φh) = dummy_op.biforms[k](xhs,uk,vk,φh) + _l(vk,φh) = dummy_op.liforms[k](xhs,vk,φh) + ∂Rk∂φ = ∇((uk,vk,φh) -> _a(uk,vk,φh) - _l(vk,φh),[xh_k,λh_k,φh],3) + xhs, ∂Rs∂φ = (xhs...,xh_k), (∂Rs∂φ...,∂Rk∂φ) + end + return ∂Rs∂φ +end + +function adjoint_solve!(φ_to_u::StaggeredAffineFEStateMap,du::AbstractVector,::Nothing) + +end + +function adjoint_solve!(φ_to_u::StaggeredAffineFEStateMap,du::AbstractVector,cache) + +end + +function pullback(φ_to_u::StaggeredAffineFEStateMap,uh,φh,du,cache) + dudφ_vec, old_adjoint_cache = cache + + ## Adjoint Solve + λ, adjoint_cache = adjoint_solve!(φ_to_u, du, old_adjoint_cache) + λh = FEFunction(get_test_space(φ_to_u),λ) + + ## Compute grad + dudφ_vecdata = collect_cell_vector(get_deriv_space(φ_to_u),dRdφ(φ_to_u,uh,λh,φh)) + assemble_vector!(dudφ_vec,get_deriv_assembler(φ_to_u),dudφ_vecdata) + rmul!(dudφ_vec, -1) + + return (NoTangent(),dudφ_vec) +end \ No newline at end of file diff --git a/scripts/Embedded/MWEs/StaggeredFEStateMaps/extensions.jl b/scripts/Embedded/MWEs/StaggeredFEStateMaps/extensions.jl new file mode 100644 index 0000000..0f920db --- /dev/null +++ b/scripts/Embedded/MWEs/StaggeredFEStateMaps/extensions.jl @@ -0,0 +1,112 @@ +struct StaggeredAdjointAffineFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB} + biforms :: Vector{<:Function} + liforms :: Vector{<:Function} + trials :: Vector{<:FESpace} + tests :: Vector{<:FESpace} + assems :: Vector{<:Assembler} + trial :: GridapSolvers.BlockSolvers.BlockFESpaceTypes{NB,SB} + test :: GridapSolvers.BlockSolvers.BlockFESpaceTypes{NB,SB} + + @doc """ + function StaggeredAdjointAffineFEOperator( + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace}, + [assems :: Vector{<:Assembler}] + ) + + Constructor for a `StaggeredAdjointAffineFEOperator` 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 StaggeredAdjointAffineFEOperator( + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trials :: Vector{<:FESpace}, + tests :: Vector{<:FESpace}, + assems :: Vector{<:Assembler} = map(SparseMatrixAssembler,trials,tests) + ) + @assert length(biforms) == length(liforms) == length(trials) == length(tests) == length(assems) + trial = combine_fespaces(trials) + test = combine_fespaces(tests) + NB, SB = length(trials), Tuple(map(num_fields,trials)) + new{NB,SB}(biforms,liforms,trials,tests,assems,trial,test) + end + + @doc """ + function StaggeredAdjointAffineFEOperator( + 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 StaggeredAdjointAffineFEOperator( + biforms :: Vector{<:Function}, + liforms :: Vector{<:Function}, + trial :: GridapSolvers.BlockSolvers.BlockFESpaceTypes{NB,SB,P}, + test :: GridapSolvers.BlockSolvers.BlockFESpaceTypes{NB,SB,P}, + assem :: GridapSolvers.BlockSolvers.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::StaggeredAdjointAffineFEOperator) = op.trial +FESpaces.get_test(op::StaggeredAdjointAffineFEOperator) = op.test + +function GridapSolvers.BlockSolvers.get_operator(op::StaggeredAdjointAffineFEOperator{NB}, xhs, k) where NB + @assert NB >= k + a(uk,vk) = op.biforms[k](xhs,uk,vk) + l(vk) = op.liforms[k](xhs,vk) + A = GridapTopOpt.assemble_adjoint_matrix(a,op.assems[k],op.trials[k],op.tests[k]) + b = assemble_vector(l,op.assems[k],op.tests[k]) + return AffineFEOperator(op.trials[k],op.tests[k],AffineOperator(A,b)) +end + +function GridapSolvers.BlockSolvers.get_operator!(op_k::AffineFEOperator, op::StaggeredAdjointAffineFEOperator{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) + assemble_matrix_and_vector!(a,l,A,b,op.assems[k],op.trials[k],op.tests[k],zero(op.tests[k])) + return op_k +end + +# Not currently required +function Algebra.solve!(xh, solver::StaggeredFESolver{NB}, op::StaggeredFEOperator{NB},l0::AbstractVector,::Nothing) where NB + solvers = solver.solvers + xhs, caches, operators = (), (), () + for k in 1:NB + xh_k = get_solution(op,xh,k) + op_k = get_operator(op,xhs,k) + l = get_vector(op_k); l .+= l0 + xh_k, cache_k = solve!(xh_k,solvers[k],op_k,nothing) + xhs, caches, operators = (xhs...,xh_k), (caches...,cache_k), (operators...,op_k) + end + return xh, (caches,operators) +end + +function Algebra.solve!(xh, solver::StaggeredFESolver{NB}, op::StaggeredFEOperator{NB}, l0::AbstractVector,cache) where NB + last_caches, last_operators = cache + solvers = solver.solvers + xhs, caches, operators = (), (), () + for k in 1:NB + xh_k = get_solution(op,xh,k) + op_k = get_operator!(last_operators[k],op,xhs,k) + l = get_vector(op_k); l .+= l0 + xh_k, cache_k = solve!(xh_k,solvers[k],op_k,last_caches[k]) + xhs, caches, operators = (xhs...,xh_k), (caches...,cache_k), (operators...,op_k) + end + return xh, (caches,operators) +end