From d5fb7284deb360b5def36a606f87a87e18775555 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Wed, 15 Jan 2025 09:59:36 +1100 Subject: [PATCH 1/8] 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 eedbee28..8b2c3a7e 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 689a0621..5fa989ca 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 69cc8ddd..bbaa8ed8 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 From 6fb8c59d5a66aadc7c9dd29cf5458ace4d3616d6 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Wed, 12 Feb 2025 09:27:38 +1100 Subject: [PATCH 2/8] Update for skeletons --- src/Autodiff.jl | 128 +++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 110 insertions(+), 18 deletions(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index 8b2c3a7e..1afe6f99 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -13,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(FESpaces._get_cell_dof_values,local_views(uh),local_trians) + cell_u = map(FESpaces.get_cell_dof_values,local_views(uh)) 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) @@ -31,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(FESpaces._get_cell_dof_values,local_views(uh),local_trians) + cell_u = map(FESpaces.get_cell_dof_values,local_views(uh)) 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) @@ -49,7 +49,7 @@ function FESpaces._hessian(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(hessian,f,local_trians,uh) - cell_u = map(FESpaces._get_cell_dof_values,local_views(uh),local_trians) + cell_u = map(FESpaces.get_cell_dof_values,local_views(uh)) 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) @@ -62,28 +62,24 @@ end # - 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) - ) + function dist_cf(uh::DistributedCellField,cfs) + DistributedCellField(cfs,get_triangulation(uh)) end - function dist_cf(uh::DistributedMultiFieldCellField,trians,spaces,cell_u) - mf_cfs = map(serial_cf,trians,spaces,cell_u) + function dist_cf(uh::DistributedMultiFieldCellField,cfs) sf_cfs = map(DistributedCellField, - [tuple_of_arrays(map(cf -> Tuple(cf.single_fields),mf_cfs))...], + [tuple_of_arrays(map(cf -> Tuple(cf.single_fields),cfs))...], map(get_triangulation,uh) ) - DistributedMultiFieldCellField(sf_cfs,mf_cfs) + DistributedMultiFieldCellField(sf_cfs,cfs) end - spaces = map(get_fe_space,local_views(uh)) + uhs = local_views(uh) + spaces = map(get_fe_space,uhs) function g(cell_u) - cf = dist_cf(uh,local_trians,spaces,cell_u) - cell_grad = f(cf) - map(get_contribution,local_views(cell_grad),local_trians) + cfs = map(CellField,spaces,cell_u) + cf = dist_cf(uh,cfs) + cg = f(cf) + map(get_contribution,local_views(cg),local_trians) end g end @@ -162,3 +158,99 @@ 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 + +# Skeleton AD + +function FESpaces._change_argument(op,f,local_trians::AbstractArray{<:SkeletonTriangulation},uh::DistributedADTypes) + function dist_cf(uh::DistributedCellField,cfs) + DistributedCellField(cfs, get_triangulation(uh)) + end + function dist_cf(uh::DistributedMultiFieldCellField,cfs) + sf_cfs = map(DistributedCellField, + [tuple_of_arrays(map(cf -> Tuple(cf.single_fields),cfs))...], + map(get_triangulation,uh) + ) + DistributedMultiFieldCellField(sf_cfs,cfs) + end + + uhs = local_views(uh) + spaces = map(get_fe_space,uhs) + function g(cell_u) + uhs_dual = map(CellField,spaces,cell_u) + cf_plus = dist_cf(uh,map(SkeletonCellFieldPair,uhs_dual,uhs)) + cf_minus = dist_cf(uh,map(SkeletonCellFieldPair,uhs,uhs_dual)) + cg_plus = f(cf_plus) + cg_minus = f(cf_minus) + plus = map(get_contribution,local_views(cg_plus),local_trians) + minus = map(get_contribution,local_views(cg_minus),local_trians) + plus, minus + end + g +end + +function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) + 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 + + # dual output of both sides at once + j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual) + + j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus + # Work for plus side + j_to_cfg_plus = lazy_map(Reindex(i_to_cfg),j_to_i.plus) + j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus) + + # Work for minus side + j_to_cfg_minus = lazy_map(Reindex(i_to_cfg),j_to_i.minus) + j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_minus) + + # Assemble on SkeletonTriangulation expects an array of interior of facets + # where each entry is a 2-block BlockVector with the first block being the + # contribution of the plus side and the second, the one of the minus side + is_single_field = eltype(eltype(j_to_result_plus)) <: Number + k = is_single_field ? BlockMap(2,[1,2]) : Fields.BlockBroadcasting(BlockMap(2,[1,2])) + lazy_map(k,j_to_result_plus,j_to_result_minus) + end + + return j_to_result +end + +function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) + 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 + + # dual output of both sides at once + j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual) + + j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus + # Work for plus side + j_to_cfg_plus = lazy_map(Reindex(i_to_cfg),j_to_i.plus) + j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus) + + # Work for minus side + j_to_cfg_minus = lazy_map(Reindex(i_to_cfg),j_to_i.minus) + j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_minus) + + # Merge the columns into a 2x2 block matrix + # I = [[(CartesianIndex(i,),CartesianIndex(i,j)) for i in 1:2] for j in 1:2] + I = [ + [(CartesianIndex(1,), CartesianIndex(1, 1)), (CartesianIndex(2,), CartesianIndex(2, 1))], # Plus -> First column + [(CartesianIndex(1,), CartesianIndex(1, 2)), (CartesianIndex(2,), CartesianIndex(2, 2))] # Minus -> Second column + ] + is_single_field = eltype(eltype(j_to_result_plus)) <: AbstractArray + k = is_single_field ? Fields.MergeBlockMap((2,2),I) : Fields.BlockBroadcasting(Fields.MergeBlockMap((2,2),I)) + lazy_map(k,j_to_result_plus,j_to_result_minus) + end + + return j_to_result +end From a03b2d90b4f9434cbfcc3b62bfec7b04699afb81 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Wed, 12 Feb 2025 09:48:29 +1100 Subject: [PATCH 3/8] Minor --- src/Autodiff.jl | 2 +- test/AutodiffTests.jl | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index 1afe6f99..5d1bf797 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -223,7 +223,7 @@ end function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) dummy_tag = ()->() i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),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) diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index bbaa8ed8..20e67951 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -4,6 +4,8 @@ using Test using Gridap, Gridap.Algebra using GridapDistributed using PartitionedArrays +using SparseArrays +using ForwardDiff function main_sf(distribute,parts) ranks = distribute(LinearIndices((prod(parts),))) @@ -40,6 +42,16 @@ function main_sf(distribute,parts) b = assemble_vector(dg,U) b_AD = assemble_vector(gradient(g,uh),U) @test b ≈ b_AD + + # Skeleton AD + # I would like to compare the results, but we cannot be using FD in parallel... + Λ = SkeletonTriangulation(model) + dΛ = Measure(Λ,2*k) + g_Λ(v) = ∫(mean(v))*dΛ + r_Λ(u,v) = ∫(mean(u)*mean(v))*dΛ + + b_Λ_AD = assemble_vector(gradient(g_Λ,uh),U) + A_Λ_AD = jacobian(FEOperator(r_Λ,U,V),uh) end function main_mf(distribute,parts) From 95633733a48b3eae143adf585eeef36b4e1177eb Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Wed, 12 Feb 2025 15:46:06 +1100 Subject: [PATCH 4/8] Added more cases to add_ghost_cells that tries to reuse the same triangulation if possible --- src/Geometry.jl | 77 ++++++++++++++++++++++++++----------------------- 1 file changed, 41 insertions(+), 36 deletions(-) diff --git a/src/Geometry.jl b/src/Geometry.jl index 026ea937..f69d003c 100644 --- a/src/Geometry.jl +++ b/src/Geometry.jl @@ -741,33 +741,39 @@ function add_ghost_cells( dmodel::DistributedDiscreteModel{Dm}, dtrian::DistributedTriangulation{Dt} ) where {Dm,Dt} - covers_all_faces = _covers_all_faces(dmodel,dtrian) - if (covers_all_faces) - trians = map(local_views(dmodel)) do model - Triangulation(ReferenceFE{Dt},model) - end - return DistributedTriangulation(trians,dmodel) - else - mcell_intrian = map(local_views(dmodel),local_views(dtrian)) do model, trian - glue = get_glue(trian,Val(Dt)) - @assert isa(glue,FaceToFaceGlue) - nmcells = num_faces(model,Dt) - mcell_intrian = fill(false,nmcells) - tcell_to_mcell = glue.tface_to_mface - mcell_intrian[tcell_to_mcell] .= true - mcell_intrian - end - gids = get_face_gids(dmodel,Dt) - cache=fetch_vector_ghost_values_cache(mcell_intrian,partition(gids)) - fetch_vector_ghost_values!(mcell_intrian,cache) |> wait + tface_to_mface = map(local_views(dtrian)) do trian + glue = get_glue(trian,Val(Dt)) + @assert isa(glue,FaceToFaceGlue) + glue.tface_to_mface + end + + # Case 1: All model faces are already in the triangulation + covers_all_faces = reduce(&,map(x -> isa(x,IdentityVector), tface_to_mface),init=true) + covers_all_faces && return dtrian - dreffes = map(local_views(dmodel)) do model - ReferenceFE{Dt} - end - trians = map(Triangulation,dreffes,local_views(dmodel),mcell_intrian) - return DistributedTriangulation(trians,dmodel) + # Otherwise: Add ghost cells to triangulation + mface_intrian = map(local_views(dmodel),tface_to_mface) do model, tface_to_mface + mface_intrian = fill(false,num_faces(model,Dt)) + mface_intrian[tface_to_mface] .= true + mface_intrian end + consistent!(PVector(mface_intrian,partition(get_face_gids(dmodel,Dt)))) |> wait + + # Case 2: New triangulation contains all model faces + covers_all_faces = reduce(&,map(all,mface_intrian),init=true) + covers_all_faces && return Triangulation(with_ghost,ReferenceFE{Dt},dmodel) + + # Case 3: Original triangulation already had the ghost cells + new_tface_to_mface = map(findall,mface_intrian) + had_ghost = reduce(&,map(==,tface_to_mface,new_tface_to_mface),init=true) + had_ghost && return dtrian + + # Case 4: Otherwise, create a new triangulation with the ghost cells + trians = map(local_views(dmodel),new_tface_to_mface) do model, new_tface_to_mface + Triangulation(ReferenceFE{Dt},model,new_tface_to_mface) + end + return DistributedTriangulation(trians,dmodel) end function generate_cell_gids(dtrian::DistributedTriangulation) @@ -778,11 +784,11 @@ end function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm}, dtrian::DistributedTriangulation{Dt}) where {Dm,Dt} + mgids = get_face_gids(dmodel,Dt) covers_all_faces = _covers_all_faces(dmodel,dtrian) if (covers_all_faces) - get_face_gids(dmodel,Dt) + tgids = mgids else - mgids = get_face_gids(dmodel,Dt) # count number owned cells notcells, tcell_to_mcell = map( local_views(dmodel),local_views(dtrian),PArrays.partition(mgids)) do model,trian,partition @@ -802,7 +808,7 @@ function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm}, # Assign global cell ids to owned cells mcell_to_gtcell = map( - first_gtcell,tcell_to_mcell,PArrays.partition(mgids)) do first_gtcell,tcell_to_mcell,partition + first_gtcell,tcell_to_mcell,partition(mgids)) do first_gtcell,tcell_to_mcell,partition mcell_to_gtcell = zeros(Int,local_length(partition)) loc_to_owner = local_to_owner(partition) part = part_id(partition) @@ -816,22 +822,21 @@ function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm}, mcell_to_gtcell end - cache = fetch_vector_ghost_values_cache(mcell_to_gtcell,PArrays.partition(mgids)) + cache = fetch_vector_ghost_values_cache(mcell_to_gtcell,partition(mgids)) fetch_vector_ghost_values!(mcell_to_gtcell,cache) |> wait # Prepare new partition ngtcells = reduction(+,notcells,destination=:all,init=zero(eltype(notcells))) - partition = map(ngtcells, - mcell_to_gtcell, - tcell_to_mcell, - PArrays.partition(mgids)) do ngtcells,mcell_to_gtcell,tcell_to_mcell,partition + indices = map( + ngtcells,mcell_to_gtcell,tcell_to_mcell,partition(mgids) + ) do ngtcells,mcell_to_gtcell,tcell_to_mcell,partition tcell_to_gtcell = mcell_to_gtcell[tcell_to_mcell] - lid_to_owner = local_to_owner(partition) + lid_to_owner = local_to_owner(partition) tcell_to_part = lid_to_owner[tcell_to_mcell] LocalIndices(ngtcells,part_id(partition),tcell_to_gtcell,tcell_to_part) end - _find_neighbours!(partition, PArrays.partition(mgids)) - gids = PRange(partition) - gids + _find_neighbours!(indices, partition(mgids)) + tgids = PRange(indices) end + return tgids end From 0384e595c3d2ec1a30cd2b962a22c4c09ad9c976 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Thu, 20 Feb 2025 12:00:01 +1100 Subject: [PATCH 5/8] Fixed num_cells for empty parts --- src/Geometry.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Geometry.jl b/src/Geometry.jl index f69d003c..57ec00a9 100644 --- a/src/Geometry.jl +++ b/src/Geometry.jl @@ -486,7 +486,7 @@ function Geometry.num_cells(a::DistributedTriangulation{Df}) where Df else mcell_to_owned = local_to_own(gids) is_owned(mcell) = !iszero(mcell_to_owned[mcell]) - sum(is_owned,tcell_to_mcell) + sum(is_owned,tcell_to_mcell;init=0) end end return sum(n_loc_ocells) From 1419e1ff56e579e07115e8f0324d17ade1c738d9 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Thu, 20 Feb 2025 14:28:38 +1100 Subject: [PATCH 6/8] Minor --- src/Autodiff.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index 5d1bf797..e045cd6f 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -132,7 +132,7 @@ function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) end j_to_ydual = a(i_to_xdual) 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) + j_to_cfg = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i) lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual) end return j_to_result @@ -148,7 +148,7 @@ function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i) end j_to_ydual = a(i_to_xdual) 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) + j_to_cfg = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i) lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual) end return j_to_result @@ -202,11 +202,11 @@ function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<: j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus # Work for plus side - j_to_cfg_plus = lazy_map(Reindex(i_to_cfg),j_to_i.plus) + j_to_cfg_plus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.plus) j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus) # Work for minus side - j_to_cfg_minus = lazy_map(Reindex(i_to_cfg),j_to_i.minus) + j_to_cfg_minus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.minus) j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_minus) # Assemble on SkeletonTriangulation expects an array of interior of facets @@ -234,11 +234,11 @@ function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<: j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus # Work for plus side - j_to_cfg_plus = lazy_map(Reindex(i_to_cfg),j_to_i.plus) + j_to_cfg_plus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.plus) j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus) # Work for minus side - j_to_cfg_minus = lazy_map(Reindex(i_to_cfg),j_to_i.minus) + j_to_cfg_minus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.minus) j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_minus) # Merge the columns into a 2x2 block matrix From ac738feaeb2f204ac1cb88e42eaff46485527b44 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Tue, 4 Mar 2025 11:41:55 +1100 Subject: [PATCH 7/8] Updated NEWS and bumped version --- NEWS.md | 6 ++++++ Project.toml | 2 +- src/Autodiff.jl | 4 ---- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/NEWS.md b/NEWS.md index 3ba5fb2a..c9c3f110 100644 --- a/NEWS.md +++ b/NEWS.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.7] 2025-03-04 + +### Added + +- Extended support for automatic differentiation to multi-field spaces and skeleton triangulations. Since PR[#169](https://github.com/gridap/GridapDistributed.jl/pull/169). + ## [0.4.6] 2024-12-03 ### Added diff --git a/Project.toml b/Project.toml index 9a551a0a..f31b3b48 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "GridapDistributed" uuid = "f9701e48-63b3-45aa-9a63-9bc6c271f355" authors = ["S. Badia ", "A. F. Martin ", "F. Verdugo "] -version = "0.4.6" +version = "0.4.7" [deps] BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" diff --git a/src/Autodiff.jl b/src/Autodiff.jl index e045cd6f..558800f2 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -57,10 +57,6 @@ function FESpaces._hessian(f,uh,fuh::DistributedDomainContribution) 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) function dist_cf(uh::DistributedCellField,cfs) DistributedCellField(cfs,get_triangulation(uh)) From 86d0d9389ed08fab737e3a7c5532c8f54f513e3c Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Tue, 4 Mar 2025 12:12:00 +1100 Subject: [PATCH 8/8] Minor --- test/TestApp/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/TestApp/Project.toml b/test/TestApp/Project.toml index 31c314de..8e3f9472 100644 --- a/test/TestApp/Project.toml +++ b/test/TestApp/Project.toml @@ -4,6 +4,7 @@ version = "0.1.0" [deps] BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e" GridapDistributed = "f9701e48-63b3-45aa-9a63-9bc6c271f355" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"