How to use Autodiff in Gridap.jl

Specific Domains Modelling & Simulations
, ,
1

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.
2

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?

3

Not sure.

What triggers the method error? Does the PDE solve in itself work fine? Can the FE assembly (thus excluding the linear solve) be auto-differentiated? Does Gridap have examples using ForwardDiff ?

An example of a FEM solve + ForwardDiff is Landau in Ferrite. Not use how relevant this is. The naive intuition is that this should be possible in Gridap as well.

4

Hey there.

So I’m trying to solve EIT (Electrical Impedance tomography) for a while now in Gridap.jl
Here is my code for that all in one file:

I failed using Autodiff on their Matrix assembler. All the Autodiff libraries (Zygote, ForwardDiff) failed.
So I used Adjoint methods
Which is easy here since the poisson operator is self adjoint.
Now the step I find most difficult is the functional derivative: ∇(u)⋅∇(λ)
Is there a good way to assemble that tensor once and reuse it from then on?

I mean I would like to have a chat with Prof. Verdugo or Badia and actually ask how one is supposed to solve such a problem with their library. And I also have some question about the internals and design decisions.

I’m stuck with my thesis since more than a year since I don’t get the FEM to work.

And I haven’t even started with using a convolutional network on the mesh with Lux.jl

Sorry that my answer isn’t helpful but I tried and it didn’t work.

5

See for example this topology optimization tutorial in Gridap. It doesn’t use AD, but manually implements an adjoint method.

(In general, AD tools seem to really struggle at differentiating through an FEM solver, unless you give them a bit of help by manually implementing vector–Jacobian products for some key parts.)