Split AD for Multifield

NEWS.md

@@ -5,11 +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).
 
-## [Unreleased]
+## [0.4.10] - 2025-09-29
 
 ### Added
 
 - Added support for multiple ghost layers on cartesian models. Since PR[#182](https://github.com/gridap/GridapDistributed.jl/pull/182).
+- Added new way of doing AD for MultiField, where partials are computed separately for each field then merged together. Since PR[#176](https://github.com/gridap/GridapDistributed.jl/pull/176).
 
 ### Fixed
 
@@ -159,7 +160,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-- Added missing parameter to `allocate_jacobian`, needed after Gridap v0.17.18. Since PR [126](https://github.com/gridap/GridapDistributed.jl/pull/126). 
+- Added missing parameter to `allocate_jacobian`, needed after Gridap v0.17.18. Since PR [126](https://github.com/gridap/GridapDistributed.jl/pull/126).
 
 ## [0.2.8] - 2023-07-31
 

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.9"
+version = "0.4.10"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

src/Autodiff.jl

@@ -80,6 +80,34 @@ function FESpaces._change_argument(op,f,local_trians,uh::DistributedADTypes)
   g
 end
 
+# Distributed counterpart of: src/MultiField/MultiFieldAutodiff.jl
+
+for (op,_op) in ((:gradient,:_gradient),(:jacobian,:_jacobian))
+  @eval begin
+    function FESpaces.$(op)(f::Function,uh::DistributedMultiFieldFEFunction;ad_type=:split)
+      fuh = f(uh)
+      if ad_type == :split
+        MultiField.multifield_autodiff_split($op,f,uh,fuh)
+      elseif ad_type == :monolithic
+        FESpaces.$(_op)(f,uh,fuh)
+      else
+        @notimplemented """Unknown ad_type = $ad_type
+          Options:
+          - :split      -- compute the gradient for each field separately, then merge
+          - :monolithic -- compute the gradient for all fields together
+          """
+      end
+    end
+  end
+end
+
+function MultiField._combine_contributions(op::Function,terms,fuh::DistributedDomainContribution)
+  local_terms = map(local_views(fuh),local_views.(terms)...) do fuh,terms...
+    MultiField._combine_contributions(op,terms,fuh)
+  end
+  DistributedDomainContribution(local_terms)
+end
+
 # Distributed counterpart of: src/Arrays/Autodiff.jl
 # autodiff_array_xxx
 

src/Geometry.jl

@@ -394,7 +394,7 @@ function _cartesian_model_with_periodic_bcs(pdesc::DistributedCartesianDescripto
     remove_boundary = map((p,n)->((p && (n!=1)) ? true : false),desc.isperiodic,parts)
     CartesianDiscreteModel(_desc,cmin,cmax,remove_boundary)
   end
-  
+
   return models, global_partition
 end
 
@@ -659,7 +659,7 @@ function Geometry.SkeletonTriangulation(
   return filter_cells_when_needed(portion,gids,parent)
 end
 
-# NOTE: The following constructors require adding back the ghost cells: 
+# NOTE: The following constructors require adding back the ghost cells:
 # Potentially, the input `trian` has had some/all of its ghost cells removed. If we do not
 # add them back, some skeleton facets might look like boundary facets to the local constructors...
 function Geometry.BoundaryTriangulation(
@@ -937,4 +937,4 @@ function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm},
     tgids = PRange(indices)
   end
   return tgids
-end
+end

src/MultiField.jl

@@ -33,6 +33,8 @@ 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.getindex(m::DistributedMultiFieldCellField,field_id::AbstractUnitRange) = m.field_fe_fun[field_id]
+Base.lastindex(m::DistributedMultiFieldCellField) = num_fields(m)
 Base.length(m::DistributedMultiFieldCellField) = num_fields(m)
 
 function LinearAlgebra.dot(a::DistributedMultiFieldCellField,b::DistributedMultiFieldCellField)
@@ -161,7 +163,7 @@ function FESpaces.EvaluationFunction(
   isconsistent=false
 )
   free_values = change_ghost(_free_values,f.gids;is_consistent=isconsistent,make_consistent=true)
-  
+
   # Create distributed single field functions
   field_fe_fun = DistributedSingleFieldFEFunction[]
   for i in 1:num_fields(f)
@@ -221,7 +223,7 @@ function FESpaces.interpolate_everywhere!(
 )
   msg = "free_values and FESpace have incompatible index partitions."
   @check PartitionedArrays.matching_local_indices(axes(free_values,1),get_free_dof_ids(space)) msg
-  
+
   # Interpolate each field
   field_fe_fun = DistributedSingleFieldFEFunction[]
   for i in 1:num_fields(space)
@@ -396,16 +398,16 @@ function generate_multi_field_gids(
     end
     collect(keys(dict))
   end
-  
+
   f_p_parts_snd, f_p_parts_rcv = map(x->assembly_neighbors(partition(x)),f_frange) |> tuple_of_arrays
   p_f_parts_snd = map(v,f_p_parts_snd...)
   p_f_parts_rcv = map(v,f_p_parts_rcv...)
   p_neigs_snd = map(merge_neigs,p_f_parts_snd)
   p_neigs_rcv = map(merge_neigs,p_f_parts_rcv)
-  
+
   exchange_graph = ExchangeGraph(p_neigs_snd,p_neigs_rcv)
   assembly_neighbors(p_iset;neighbors=exchange_graph)
-  
+
   PRange(p_iset)
 end
 
@@ -446,7 +448,7 @@ end
 
 # BlockSparseMatrixAssemblers
 
-const DistributedBlockSparseMatrixAssembler{R,C} = 
+const DistributedBlockSparseMatrixAssembler{R,C} =
   MultiField.BlockSparseMatrixAssembler{R,C,<:AbstractMatrix{<:DistributedSparseMatrixAssembler}}
 
 function FESpaces.SparseMatrixAssembler(

test/AutodiffTests.jl

@@ -44,7 +44,7 @@ function main_sf(distribute,parts)
   @test b ≈ b_AD
 
   # Skeleton AD
-  # I would like to compare the results, but we cannot be using FD in parallel... 
+  # 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Λ
@@ -73,7 +73,7 @@ function main_mf(distribute,parts)
 
   Ω = Triangulation(model)
   dΩ = Measure(Ω,2*(k+1))
-  
+
   ν = 1.0
   f = VectorValue(0.0,0.0)
 
@@ -104,9 +104,113 @@ function main_mf(distribute,parts)
   @test b ≈ b_AD
 end
 
+## MultiField AD with different triangulations for each field
+function generate_trian(ranks,model,case)
+  cell_ids = get_cell_gids(model)
+  trians = map(ranks,local_views(model),partition(cell_ids)) do rank, model, ids
+    cell_mask = zeros(Bool, num_cells(model))
+    if case == :partial_trian
+      if rank ∈ (1,2)
+        cell_mask[own_to_local(ids)] .= true
+      else
+        t = own_to_local(ids)
+        cell_mask[t[1:floor(Int,length(t)/2)]] .= true
+      end
+    elseif case == :half_empty_trian
+      if rank ∈ (3,4)
+        cell_mask[own_to_local(ids)] .= true
+      end
+    elseif case == :trian_with_empty_procs
+      if rank ∈ (1,2)
+        t = own_to_local(ids)
+        cell_mask[t[1:floor(Int,length(t)/2)]] .= true
+      end
+    else
+      error("Unknown case")
+    end
+    Triangulation(model,cell_mask)
+  end
+  GridapDistributed.DistributedTriangulation(trians,model)
+end
+
+function mf_different_fespace_trians(distribute,parts)
+  ranks = distribute(LinearIndices((prod(parts),)))
+  model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(10,10))
+  V2 = FESpace(model,ReferenceFE(lagrangian,VectorValue{2,Float64},1))
+  V3 = FESpace(model,ReferenceFE(lagrangian,Float64,1))
+  for case in (:boundary,:partial_trian, :half_empty_trian, :trian_with_empty_procs)
+    if case == :boundary
+      Γ = BoundaryTriangulation(model)
+    else
+      Γ = generate_trian(ranks,model,case)
+    end
+    dΓ = Measure(Γ,2)
+    V1 = FESpace(Γ,ReferenceFE(lagrangian,Float64,1))
+    X = MultiFieldFESpace([V1,V2,V3])
+    uh = zero(X);
+
+    f(xh) = ∫(xh[1]+xh[2]⋅xh[2]+xh[1]*xh[3])dΓ
+    df(v,xh) = ∫(v[1]+2*v[2]⋅xh[2]+v[1]*xh[3]+xh[1]*v[3])dΓ
+    du = gradient(f,uh)
+    du_vec = assemble_vector(du,X)
+    df_vec = assemble_vector(v->df(v,uh),X)
+
+    @test df_vec ≈ du_vec
+
+    f2(xh,yh) = ∫(xh[1]⋅yh[1]+xh[2]⋅yh[2]+xh[1]⋅xh[2]⋅yh[2]+xh[1]*xh[3]*yh[3])dΓ
+    dv = get_fe_basis(X);
+    j = jacobian(uh->f2(uh,dv),uh)
+    J = assemble_matrix(j,X,X)
+
+    f2_jac(xh,dxh,yh) = ∫(dxh[1]⋅yh[1]+dxh[2]⋅yh[2]+dxh[1]⋅xh[2]⋅yh[2]+xh[1]⋅dxh[2]⋅yh[2]+dxh[1]*xh[3]*yh[3]+xh[1]*dxh[3]*yh[3])dΓ
+    op = FEOperator(f2,f2_jac,X,X)
+    J_fwd = jacobian(op,uh)
+
+    @test reduce(&,map(≈,partition(J),partition(J_fwd)))
+  end
+end
+
+function skeleton_mf_different_fespace_trians(distribute,parts)
+  ranks = distribute(LinearIndices((prod(parts),)))
+  model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(10,10))
+  for case in (:partial_trian, :half_empty_trian, :trian_with_empty_procs)
+    Γ = generate_trian(ranks,model,case)
+    V1 = FESpace(Γ,ReferenceFE(lagrangian,Float64,1),conformity=:L2)
+    V2 = FESpace(model,ReferenceFE(lagrangian,VectorValue{2,Float64},1),conformity=:L2)
+    V3 = FESpace(model,ReferenceFE(lagrangian,Float64,1),conformity=:L2)
+    X = MultiFieldFESpace([V1,V2,V3])
+    uh = zero(X);
+    Λ = SkeletonTriangulation(model)
+    dΛ = Measure(Λ,2)
+
+    f(xh) = ∫(mean(xh[1])+mean(xh[2])⋅mean(xh[2])+mean(xh[1])*mean(xh[3]))dΛ
+    df(v,xh) = ∫(mean(v[1])+2*mean(v[2])⋅mean(xh[2])+mean(v[1])*mean(xh[3])+mean(xh[1])*mean(v[3]))dΛ
+    du = gradient(f,uh)
+    du_vec = assemble_vector(du,X)
+    df_vec = assemble_vector(v->df(v,uh),X)
+
+    @test df_vec ≈ du_vec
+
+    # Skel jac
+    f2(xh,yh) = ∫(mean(xh[1])⋅mean(yh[1])+mean(xh[2])⋅mean(yh[2])+mean(xh[1])⋅mean(xh[2])⋅mean(yh[2])+mean(xh[1])*mean(xh[3])*mean(yh[3]))dΛ
+    dv = get_fe_basis(X);
+    j = jacobian(uh->f2(uh,dv),uh);
+    J = assemble_matrix(j,X,X)
+
+    f2_jac(xh,dxh,yh) = ∫(mean(dxh[1])⋅mean(yh[1])+mean(dxh[2])⋅mean(yh[2])+mean(dxh[1])⋅mean(xh[2])⋅mean(yh[2]) +
+      mean(xh[1])⋅mean(dxh[2])⋅mean(yh[2])+mean(dxh[1])*mean(xh[3])*mean(yh[3])+mean(xh[1])*mean(dxh[3])*mean(yh[3]))dΛ
+    op = FEOperator(f2,f2_jac,X,X)
+    J_fwd = jacobian(op,uh)
+
+    @test reduce(&,map(≈,partition(J),partition(J_fwd)))
+  end
+end
+
 function main(distribute,parts)
   main_sf(distribute,parts)
   main_mf(distribute,parts)
+  mf_different_fespace_trians(distribute,parts)
+  skeleton_mf_different_fespace_trians(distribute,parts)
 end
 
 end