Merge pull request #169 from gridap/autodiff

AD for multifield and skeletons

Author: JordiManyer

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

Project.toml

@@ -1,7 +1,7 @@
 name = "GridapDistributed"
 uuid = "f9701e48-63b3-45aa-9a63-9bc6c271f355"
 authors = ["S. Badia <santiago.badia@monash.edu>", "A. F. Martin <alberto.f.martin@anu.edu.au>", "F. Verdugo <f.verdugo.rojano@vu.nl>"]
-version = "0.4.6"
+version = "0.4.7"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

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))
     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,222 @@ 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))
     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)
-  spaces = map(get_fe_space,local_views(uh))
-  function g(cell_u)
-    cf = DistributedCellField(
-      map(CellField,spaces,cell_u),
-      get_triangulation(uh),
+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))
+    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
+
+function FESpaces._change_argument(op,f,local_trians,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)
     )
-    cell_grad = f(cf)
-    map(get_contribution,local_views(cell_grad),trian)
+    DistributedMultiFieldCellField(sf_cfs,cfs)
+  end
+
+  uhs = local_views(uh)
+  spaces = map(get_fe_space,uhs)
+  function g(cell_u)
+    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
 
+# 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 = Arrays.autodiff_array_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 = Arrays.autodiff_array_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
+
+# 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 = 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 = 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
+    # 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.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
+
+  # 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 = 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 = 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
+    # 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