How to use Autodiff in Gridap.jl

TLDR: I want to solve an optimal control problem using automatic differentiation (Gridap uses ForwardDiff.jl), but i get an error when i try to find the gradient of the objective function. The objective function depends on a PDE-solver from Gridap.jl. Can the gradient be found by using another Autodiff-package instead? Or is there some error in the way I have defined the objective function?

I want to solve an optimal control problem, with respect to a PDE (the Poisson equation). The optimization problem is given by,

min J(y, u) =0.5 * ∫(y - y_Ω)^2 dΩ) + 0.5 * ∫(u)^2 dΩ

such that,

Δy = u  in Ω

I.e. i want to find the force function u, such that the corresponding solution of the Poisson equation minimizes the integral of the squared difference y - y_Ω.

The reduced functional f is defined as

 f(u) = J(S(u), u)

Where S is the solution operator of the Poisson equation.

However, instead of using the adjoint method, where the gradient of the functional f is computed analytically by solving the corresponding adjoint problem, I want to find the gradient using automatic differentiation.

However, it doesn’t seem that ForwardDiff is able to differentiate through the PDE solver. Can this be done in another way? E.g. by using another autodiff package? Or have I misunderstood the use of DomainContribution in the definition of f?

This is an example,

import Base: sin, π

import Gridap: Point, evaluate, interpolate, writevtk
using Gridap
using Gridap.FESpaces
using Gridap.CellData
using Gridap.Arrays

using Plots


domain = (0, 1)
partition = (8,)
order = 1

model = CartesianDiscreteModel(domain, partition)
reffe = ReferenceFE(lagrangian, Float64, order)

V0 = TestFESpace(model, reffe; conformity=:H1, dirichlet_tags="boundary")

# Boundary value function
g(x) = 0
Ug = TrialFESpace(V0, g)

# Numerical integration
Ω = Triangulation(model)


degree = 2
dΩ = Measure(Ω, degree)

inner_product(a, b) = sum(∫(a * b) * dΩ)
l2_norm(a) = sqrt(inner_product(a, a))


function solve_poisson(f)
    a(u, v) = ∫(∇(v) ⋅ ∇(u)) *dΩ
    b(v) = ∫(v * f) * dΩ

    # FE Problem
    op = AffineFEOperator(a, b, Ug, V0)

    # Solve phase
    ls = LUSolver()
    solver = LinearFESolver(ls)

    uh = solve(solver, op)
    return uh
end


# The 'goal' function we want to find a control function u(x) for
y_Ω(x) = sin(2 * π * x[1])
y_Ω_fe = interpolate(y_Ω, V0)


# Reduced cost functional, defined from the objective function by:
#   f(u) = J(Su, u) = J(y, u)
function f(u)
    y_diff = solve_poisson(u) - y_Ω_fe
    f_1 = 0.5 * inner_product(y_diff, y_diff)
    y = f_1 + 1. / 2.0 * inner_product(u, u)

    dc = CellData.DomainContribution()
    CellData.add_contribution!(dc, Ω, [y])
    return dc
end


# initial control function for the iteration
u₀(x) = 1.0
u0 = interpolate(u₀, V0)

# compute the gradient of the functional f at u0
∇(f, u0)

When running this i get the following stack trace:

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{Gridap.Arrays.var"#93#94", Float64}, Float64, 2})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:792
  Float64(::IrrationalConstants.Twoinvπ)
   @ IrrationalConstants ~/.julia/packages/IrrationalConstants/lWTip/src/macro.jl:131
  ...

Stacktrace:
  [1] convert(::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{Gridap.Arrays.var"#93#94", Float64}, Float64, 2})
    @ Base ./number.jl:7
  [2] setindex!(A::Vector{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{…}, Float64, 2}, i1::Int64)
    @ Base ./array.jl:1021
  [3] setindex!
    @ ./multidimensional.jl:698 [inlined]
  [4] add_entry!
    @ ~/.julia/packages/Gridap/aikpb/src/Algebra/AlgebraInterfaces.jl:130 [inlined]
  [5] add_entry!
    @ ~/.julia/packages/Gridap/aikpb/src/Algebra/AlgebraInterfaces.jl:119 [inlined]
  [6] _add_entries!
    @ ~/.julia/packages/Gridap/aikpb/src/Algebra/AlgebraInterfaces.jl:206 [inlined]
  [7] add_entries!
    @ ~/.julia/packages/Gridap/aikpb/src/Algebra/AlgebraInterfaces.jl:190 [inlined]
  [8] evaluate!
    @ ~/.julia/packages/Gridap/aikpb/src/Arrays/AlgebraMaps.jl:51 [inlined]
  [9] _numeric_loop_matvec!(A::Gridap.Algebra.InserterCSC{…}, b::Vector{…}, caches::Tuple{…}, cell_vals::LazyArray{…}, cell_rows::Table{…}, cell_cols::Table{…})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/SparseMatrixAssemblers.jl:368
 [10] numeric_loop_matrix_and_vector!(A::Gridap.Algebra.InserterCSC{…}, b::Vector{…}, a::GenericSparseMatrixAssembler, data::Tuple{…})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/SparseMatrixAssemblers.jl:351
 [11] assemble_matrix_and_vector(a::GenericSparseMatrixAssembler, data::Tuple{Tuple{…}, Tuple{…}, Tuple{…}})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/SparseMatrixAssemblers.jl:101
 [12] AffineFEOperator(weakform::Gridap.FESpaces.var"#60#61"{…}, trial::TrialFESpace{…}, test::UnconstrainedFESpace{…}, assem::GenericSparseMatrixAssembler)
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/AffineFEOperators.jl:38
 [13] AffineFEOperator(::Function, ::TrialFESpace{UnconstrainedFESpace{…}}, ::Vararg{Any})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/AffineFEOperators.jl:47
 [14] AffineFEOperator(::var"#a#11", ::var"#b#12"{…}, ::TrialFESpace{…}, ::Vararg{…})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/AffineFEOperators.jl:51
 [15] solve_poisson(f::GenericCellField{ReferenceDomain})
    @ Main ~/git/personal/fem-div/OptimalControl.jl/examples/elliptic_ad2.jl:42
 [16] f(u::GenericCellField{ReferenceDomain})
    @ Main ~/git/personal/fem-div/OptimalControl.jl/examples/elliptic_ad2.jl:61
 [17] (::Gridap.FESpaces.var"#g#65"{…})(cell_u::LazyArray{…})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/FEAutodiff.jl:102
 [18] autodiff_array_gradient(a::Gridap.FESpaces.var"#g#65"{…}, i_to_x::LazyArray{…}, j_to_i::Vector{…})
    @ Gridap.Arrays ~/.julia/packages/Gridap/aikpb/src/Arrays/Autodiff.jl:28
 [19] _gradient(f::Function, uh::SingleFieldFEFunction{GenericCellField{ReferenceDomain}}, fuh::DomainContribution)
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/FEAutodiff.jl:23
 [20] gradient(f::typeof(f), uh::SingleFieldFEFunction{GenericCellField{ReferenceDomain}})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/aikpb/src/FESpaces/FEAutodiff.jl:5
 [21] top-level scope
    @ ~/git/personal/fem-div/OptimalControl.jl/examples/elliptic_ad2.jl:76
Some type information was truncated. Use `show(err)` to see complete types.

I don’t know about the specifics of your use case but as far as autodiff packages are concerned, can you try Enzyme.jl or Mooncake.jl?