Autodiff

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).
 
+## [Unreleased]
+
+### Added
+
+- Added support for automatic differentiation with ForwardDiff. Since PR[#167](https://github.com/gridap/GridapDistributed.jl/pull/167).
+
 ## [0.4.5] 2024-10-08
 
 ### Fixed

Project.toml

@@ -6,6 +6,7 @@ version = "0.4.5"
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
@@ -17,6 +18,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 [compat]
 BlockArrays = "1"
 FillArrays = "1"
+ForwardDiff = "0.10"
 Gridap = "0.18.7"
 LinearAlgebra = "1"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"

src/Autodiff.jl

@@ -0,0 +1,105 @@
+
+function Fields.gradient(f::Function,uh::DistributedCellField)
+  fuh = f(uh)
+  FESpaces._gradient(f,uh,fuh)
+end
+
+function FESpaces._gradient(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(gradient,f,local_trians,uh)
+    cell_u = map(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)
+  fuh = f(uh)
+  FESpaces._jacobian(f,uh,fuh)
+end
+
+function FESpaces._jacobian(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(jacobian,f,local_trians,uh)
+    cell_u = map(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),
+    )
+    cell_grad = f(cf)
+    map(get_contribution,local_views(cell_grad),trian)
+  end
+  g
+end
+
+# 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)
+  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)
+  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)
+  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)
+  end
+  i_to_result
+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)
+  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)
+  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)
+  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)
+  end
+  j_to_result
+end

src/GridapDistributed.jl

@@ -24,6 +24,7 @@ using WriteVTK
 using FillArrays
 using BlockArrays
 using LinearAlgebra
+using ForwardDiff
 
 import Gridap.TensorValues: inner, outer, double_contraction, symmetric_part
 import LinearAlgebra: det, tr, cross, dot, ⋅, diag
@@ -62,4 +63,6 @@ include("ODEs.jl")
 
 include("Adaptivity.jl")
 
+include("Autodiff.jl")
+
 end # module

test/AutodiffTests.jl

@@ -0,0 +1,45 @@
+module AutodiffTests
+
+using Test
+using Gridap, Gridap.Algebra
+using GridapDistributed
+using PartitionedArrays
+
+function main(distribute,parts)
+  ranks = distribute(LinearIndices((prod(parts),)))
+
+  domain = (0,4,0,4)
+  cells = (4,4)
+  model = CartesianDiscreteModel(ranks,parts,domain,cells)
+
+  u((x,y)) = (x+y)^k
+  σ(∇u) = (1.0+∇u⋅∇u)*∇u
+  dσ(∇du,∇u) = (2*∇u⋅∇du)*∇u + (1.0+∇u⋅∇u)*∇du
+  f(x) = -divergence(y->σ(∇(u,y)),x)
+
+  k = 1
+  reffe = ReferenceFE(lagrangian,Float64,k)
+  V = TestFESpace(model,reffe,dirichlet_tags="boundary")
+  U = TrialFESpace(u,V)
+
+  Ω = Triangulation(model)
+  dΩ = Measure(Ω,2*k)
+  r(u,v) = ∫( ∇(v)⋅(σ∘∇(u)) - v*f )dΩ
+  j(u,du,v) = ∫( ∇(v)⋅(dσ∘(∇(du),∇(u))) )dΩ
+
+  op = FEOperator(r,j,U,V)
+  op_AD = FEOperator(r,U,V)
+
+  uh = interpolate(1.0,U)
+  A = jacobian(op,uh)
+  A_AD = jacobian(op_AD,uh)
+  @test reduce(&,map(≈,partition(A),partition(A_AD)))
+
+  g(v) = ∫(0.5*v⋅v)dΩ
+  dg(v) = ∫(uh⋅v)dΩ
+  b = assemble_vector(dg,U)
+  b_AD = assemble_vector(gradient(g,uh),U)
+  @test b ≈ b_AD
+end
+
+end

test/PLaplacianTests.jl

@@ -8,21 +8,20 @@ using Test
 using SparseArrays
 
 function main(distribute,parts)
-  main(distribute,parts,FullyAssembledRows(),SparseMatrixCSR{0,Float64,Int})
-  main(distribute,parts,SubAssembledRows(),SparseMatrixCSC{Float64,Int})
+  main(distribute,parts,FullyAssembledRows(),SparseMatrixCSR{0,Float64,Int},false)
+  main(distribute,parts,SubAssembledRows(),SparseMatrixCSC{Float64,Int},true)
 end
 
-function main(distribute,parts,strategy,local_matrix_type)
+function main(distribute,parts,strategy,local_matrix_type,autodiff)
   ranks  = distribute(LinearIndices((prod(parts),)))
-  output = mkpath(joinpath(@__DIR__,"output"))
 
   domain = (0,4,0,4)
   cells = (4,4)
   model = CartesianDiscreteModel(ranks,parts,domain,cells)
 
   k = 1
   u((x,y)) = (x+y)^k
-  σ(∇u) =(1.0+∇u⋅∇u)*∇u
+  σ(∇u) = (1.0+∇u⋅∇u)*∇u
   dσ(∇du,∇u) = (2*∇u⋅∇du)*∇u + (1.0+∇u⋅∇u)*∇du
   f(x) = -divergence(y->σ(∇(u,y)),x)
 
@@ -35,8 +34,12 @@ function main(distribute,parts,strategy,local_matrix_type)
   V = TestFESpace(model,reffe,dirichlet_tags="boundary")
   U = TrialFESpace(u,V)
 
-  assem=SparseMatrixAssembler(local_matrix_type,Vector{Float64},U,V,strategy)
-  op = FEOperator(r,j,U,V,assem)
+  assem = SparseMatrixAssembler(local_matrix_type,Vector{Float64},U,V,strategy)
+  if !autodiff
+    op = FEOperator(r,j,U,V,assem)
+  else
+    op = FEOperator(r,U,V,assem)
+  end
 
   uh = zero(U)
   b,A = residual_and_jacobian(op,uh)
@@ -59,7 +62,6 @@ function main(distribute,parts,strategy,local_matrix_type)
   dΩo = Measure(Ωo,2*k)
   eh = u - uh
   @test sqrt(sum(∫( abs2(eh) )dΩo)) < 1.0e-9
-
 end
 
 end # module

test/TestApp/src/TestApp.jl

@@ -18,4 +18,5 @@ module TestApp
   include("../../AdaptivityMultiFieldTests.jl")
   include("../../BlockSparseMatrixAssemblersTests.jl")
   include("../../VisualizationTests.jl")
+  include("../../AutodiffTests.jl")
 end

test/mpi/runtests_np4_body.jl

@@ -59,5 +59,8 @@ function all_tests(distribute,parts)
     PArrays.toc!(t,"Visualization")
   end
 
+  TestApp.AutodiffTests.main(distribute,parts)
+  PArrays.toc!(t,"Autodiff")
+
   display(t)
 end

test/sequential/AutodiffTests.jl

@@ -0,0 +1,7 @@
+module AutodiffTestsSeq
+using PartitionedArrays
+include("../AutodiffTests.jl")
+with_debug() do distribute
+  AutodiffTests.main(distribute,(2,2))
+end
+end # module