AD for multifield and skeletons

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,15 +13,15 @@ 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)
   end
   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,79 +31,134 @@ 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)
   end
   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

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)

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