From d5fb7284deb360b5def36a606f87a87e18775555 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Wed, 15 Jan 2025 09:59:36 +1100 Subject: [PATCH] Added AD for multifield --- src/Autodiff.jl | 135 ++++++++++++++++++++++++++++++------------ src/MultiField.jl | 1 + test/AutodiffTests.jl | 57 +++++++++++++++++- 3 files changed, 154 insertions(+), 39 deletions(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index eedbee2..8b2c3a7 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -1,5 +1,9 @@ -function Fields.gradient(f::Function,uh::DistributedCellField) +const DistributedADTypes = Union{DistributedCellField,DistributedMultiFieldCellField} + +# Distributed counterpart of: src/FESpaces/FEAutodiff.jl + +function Fields.gradient(f::Function,uh::DistributedADTypes) fuh = f(uh) FESpaces._gradient(f,uh,fuh) end @@ -9,7 +13,7 @@ function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution) local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) for local_trians in local_domains g = FESpaces._change_argument(gradient,f,local_trians,uh) - cell_u = map(get_cell_dof_values,local_views(uh)) + cell_u = map(FESpaces._get_cell_dof_values,local_views(uh),local_trians) cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians) cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id) map(add_contribution!,local_terms,local_trians,cell_grad) @@ -17,7 +21,7 @@ function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution) DistributedDomainContribution(local_terms) end -function Fields.jacobian(f::Function,uh::DistributedCellField) +function Fields.jacobian(f::Function,uh::DistributedADTypes) fuh = f(uh) FESpaces._jacobian(f,uh,fuh) end @@ -27,7 +31,7 @@ function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution) local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) for local_trians in local_domains g = FESpaces._change_argument(jacobian,f,local_trians,uh) - cell_u = map(get_cell_dof_values,local_views(uh)) + cell_u = map(FESpaces._get_cell_dof_values,local_views(uh),local_trians) cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians) cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id) map(add_contribution!,local_terms,local_trians,cell_grad) @@ -35,71 +39,126 @@ function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution) DistributedDomainContribution(local_terms) end -function FESpaces._change_argument(op,f,trian,uh::DistributedCellField) +function Fields.hessian(f::Function,uh::DistributedADTypes) + fuh = f(uh) + FESpaces._hessian(f,uh,fuh) +end + +function FESpaces._hessian(f,uh,fuh::DistributedDomainContribution) + local_terms = map(r -> DomainContribution(), get_parts(fuh)) + local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) + for local_trians in local_domains + g = FESpaces._change_argument(hessian,f,local_trians,uh) + cell_u = map(FESpaces._get_cell_dof_values,local_views(uh),local_trians) + cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians) + cell_grad = distributed_autodiff_array_hessian(g,cell_u,cell_id) + map(add_contribution!,local_terms,local_trians,cell_grad) + end + DistributedDomainContribution(local_terms) +end + +# There are 4 = 2x2 combinations, coming from: +# - Creation of the serial CellFields are different for Triangulation and SkeletonTriangulation +# - Creation of the distributed CellFields are different for DistributedCellField and DistributedMultiFieldCellField +# The internal functions take care of all 4 combinations +function FESpaces._change_argument(op,f,local_trians,uh::DistributedADTypes) + serial_cf(::Triangulation,space,cell_u) = CellField(space,cell_u) + serial_cf(::SkeletonTriangulation,space,cell_u) = SkeletonCellFieldPair(space,cell_u) + function dist_cf(uh::DistributedCellField,trians,spaces,cell_u) + DistributedCellField( + map(serial_cf,trians,spaces,cell_u), + get_triangulation(uh) + ) + end + function dist_cf(uh::DistributedMultiFieldCellField,trians,spaces,cell_u) + mf_cfs = map(serial_cf,trians,spaces,cell_u) + sf_cfs = map(DistributedCellField, + [tuple_of_arrays(map(cf -> Tuple(cf.single_fields),mf_cfs))...], + map(get_triangulation,uh) + ) + DistributedMultiFieldCellField(sf_cfs,mf_cfs) + end + spaces = map(get_fe_space,local_views(uh)) function g(cell_u) - cf = DistributedCellField( - map(CellField,spaces,cell_u), - get_triangulation(uh), - ) + cf = dist_cf(uh,local_trians,spaces,cell_u) cell_grad = f(cf) - map(get_contribution,local_views(cell_grad),trian) + map(get_contribution,local_views(cell_grad),local_trians) end g end +# Distributed counterpart of: src/Arrays/Autodiff.jl # autodiff_array_xxx function distributed_autodiff_array_gradient(a,i_to_x) - dummy_forwarddiff_tag = ()->() - i_to_xdual = map(i_to_x) do i_to_x - lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) + dummy_tag = ()->() + i_to_cfg = map(i_to_x) do i_to_x + lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x) + end + i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x + lazy_map(DualizeMap(),i_to_cfg,i_to_x) end i_to_ydual = a(i_to_xdual) - i_to_result = map(i_to_ydual,i_to_x) do i_to_ydual,i_to_x - i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) - lazy_map(AutoDiffMap(ForwardDiff.gradient),i_to_ydual,i_to_x,i_to_cfg) + i_to_result = map(i_to_cfg,i_to_ydual) do i_to_cfg,i_to_ydual + lazy_map(AutoDiffMap(),i_to_cfg,i_to_ydual) end return i_to_result end function distributed_autodiff_array_jacobian(a,i_to_x) - dummy_forwarddiff_tag = ()->() - i_to_xdual = map(i_to_x) do i_to_x - lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) + dummy_tag = ()->() + i_to_cfg = map(i_to_x) do i_to_x + lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x) + end + i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x + lazy_map(DualizeMap(),i_to_cfg,i_to_x) end i_to_ydual = a(i_to_xdual) - i_to_result = map(i_to_ydual,i_to_x) do i_to_ydual,i_to_x - i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) - lazy_map(AutoDiffMap(ForwardDiff.jacobian),i_to_ydual,i_to_x,i_to_cfg) + i_to_result = map(i_to_cfg,i_to_ydual) do i_to_cfg,i_to_ydual + lazy_map(AutoDiffMap(),i_to_cfg,i_to_ydual) end - i_to_result + return i_to_result +end + +function distributed_autodiff_array_hessian(a,i_to_x) + agrad = i_to_y -> distributed_autodiff_array_gradient(a,i_to_y) + distributed_autodiff_array_jacobian(agrad,i_to_x) end function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) - dummy_forwarddiff_tag = ()->() - i_to_xdual = map(i_to_x) do i_to_x - lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) + dummy_tag = ()->() + i_to_cfg = map(i_to_x) do i_to_x + lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x) + end + i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x + lazy_map(DualizeMap(),i_to_cfg,i_to_x) end j_to_ydual = a(i_to_xdual) - j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual - j_to_x = lazy_map(Reindex(i_to_x),j_to_i) - j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x) - lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg) + j_to_result = map(i_to_cfg,j_to_i,j_to_ydual) do i_to_cfg,j_to_i,j_to_ydual + j_to_cfg = lazy_map(Reindex(i_to_cfg),j_to_i) + lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual) end return j_to_result end function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i) - dummy_forwarddiff_tag = ()->() - i_to_xdual = map(i_to_x) do i_to_x - lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) + dummy_tag = ()->() + i_to_cfg = map(i_to_x) do i_to_x + lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x) + end + i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x + lazy_map(DualizeMap(),i_to_cfg,i_to_x) end j_to_ydual = a(i_to_xdual) - j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual - j_to_x = lazy_map(Reindex(i_to_x),j_to_i) - j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x) - lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg) + j_to_result = map(i_to_cfg,j_to_i,j_to_ydual) do i_to_cfg,j_to_i,j_to_ydual + j_to_cfg = lazy_map(Reindex(i_to_cfg),j_to_i) + lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual) end - j_to_result + return j_to_result +end + +function distributed_autodiff_array_hessian(a,i_to_x,i_to_j) + agrad = i_to_y -> distributed_autodiff_array_gradient(a,i_to_y,i_to_j) + distributed_autodiff_array_jacobian(agrad,i_to_x,i_to_j) end diff --git a/src/MultiField.jl b/src/MultiField.jl index 689a062..5fa989c 100644 --- a/src/MultiField.jl +++ b/src/MultiField.jl @@ -33,6 +33,7 @@ MultiField.num_fields(m::DistributedMultiFieldCellField) = length(m.field_fe_fun Base.iterate(m::DistributedMultiFieldCellField) = iterate(m.field_fe_fun) Base.iterate(m::DistributedMultiFieldCellField,state) = iterate(m.field_fe_fun,state) Base.getindex(m::DistributedMultiFieldCellField,field_id::Integer) = m.field_fe_fun[field_id] +Base.length(m::DistributedMultiFieldCellField) = num_fields(m) function LinearAlgebra.dot(a::DistributedMultiFieldCellField,b::DistributedMultiFieldCellField) @check num_fields(a) == num_fields(b) diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index 69cc8dd..bbaa8ed 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -5,7 +5,7 @@ using Gridap, Gridap.Algebra using GridapDistributed using PartitionedArrays -function main(distribute,parts) +function main_sf(distribute,parts) ranks = distribute(LinearIndices((prod(parts),))) domain = (0,4,0,4) @@ -42,4 +42,59 @@ function main(distribute,parts) @test b ≈ b_AD end +function main_mf(distribute,parts) + ranks = distribute(LinearIndices((prod(parts),))) + + model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(4,4)) + + k = 2 + reffe_u = ReferenceFE(lagrangian,VectorValue{2,Float64},k) + reffe_p = ReferenceFE(lagrangian,Float64,k-1;space=:P) + + u(x) = VectorValue(x[2],-x[1]) + V = TestFESpace(model,reffe_u,dirichlet_tags="boundary") + U = TrialFESpace(V,u) + Q = TestFESpace(model,reffe_p;conformity=:L2,constraint=:zeromean) + + X = MultiFieldFESpace([U,Q]) + Y = MultiFieldFESpace([V,Q]) + + Ω = Triangulation(model) + dΩ = Measure(Ω,2*(k+1)) + + ν = 1.0 + f = VectorValue(0.0,0.0) + + conv(u,∇u) = (∇u')⋅u + dconv(du,∇du,u,∇u) = conv(u,∇du)+conv(du,∇u) + c(u,v) = ∫(v⊙(conv∘(u,∇(u))))dΩ + dc(u,du,dv) = ∫(dv⊙(dconv∘(du,∇(du),u,∇(u))))dΩ + + biform((du,dp),(dv,dq)) = ∫(ν*∇(dv)⊙∇(du) - (∇⋅dv)*dp - (∇⋅du)*dq)dΩ + liform((dv,dq)) = ∫(dv⋅f)dΩ + + r((u,p),(v,q)) = biform((u,p),(v,q)) + c(u,v) - liform((v,q)) + j((u,p),(du,dp),(dv,dq)) = biform((du,dp),(dv,dq)) + dc(u,du,dv) + + op = FEOperator(r,j,X,Y) + op_AD = FEOperator(r,X,Y) + + xh = interpolate([VectorValue(1.0,1.0),1.0],X) + uh, ph = xh + A = jacobian(op,xh) + A_AD = jacobian(op_AD,xh) + @test reduce(&,map(≈,partition(A),partition(A_AD))) + + g((v,q)) = ∫(0.5*v⋅v + 0.5*q*q)dΩ + dg((v,q)) = ∫(uh⋅v + ph*q)dΩ + b = assemble_vector(dg,X) + b_AD = assemble_vector(gradient(g,xh),X) + @test b ≈ b_AD +end + +function main(distribute,parts) + main_sf(distribute,parts) + main_mf(distribute,parts) +end + end \ No newline at end of file