Zygote & GridapTopOpt Compatability

As of v0.3.0, Zygote can be used for backwards AD in GridapTopOpt in serial and distributed.

GridapTopOpt.CustomPDEConstrainedFunctionals — Type

CustomPDEConstrainedFunctionals{N,A} <:  AbstractPDEConstrainedFunctionals{N}

A version of PDEConstrainedFunctionals that allows for an arbitrary mapping φ_to_jc that is used to compute the objective and constraints given the primal variable.

Under the hood, we use Zygote to compute the Jacobian of this mapping with the rrules defined throughout GridapTopOpt.

Parameters

φ_to_jc: A function that defines the mapping from the primal variable φ to the objective and constraints. This should accept AbstractVector φ (the free values of φh) and output a scalar. In side this function, you should compute your states and evaluate your objectives and constraints that should be written as GridapTopOpt.StateParamMap. For example,
```
using GridapTopOpt: StateParamMap
...
J = StateParamMap(j,state_map)
C = StateParamMap(c,state_map)
function φ_to_jc(φ)
  u = state_map(φ)
  [J(u,φ),C(u,φ)^2]
end
pcfs = CustomPDEConstrainedFunctionals(φ_to_jc,state_map,φh)
```
analytic_dJ: Either a Function (if using analytic derivative) or nothing if using AD.
analytic_dC: A vector of Nothing or Function depending on whether using analytic derivatives or AD.
state_map::A: The state map for the problem. NOTE: this is a place holder for the optimiser output and in theory you could use many different state maps inside φ_to_jc.

Warning

The expected function for analytic_dJ and functions of analytic_dC are different to usual. Here, you should define a function that takes an AbstractVector input corresponding to the derivative and the primal variable dofs φ and assembles the derivative into the AbstractVector input. For example,

function analytic_dJ!(dJ,φ)
  φh = FEFunction(V_φ,φ)
  uh = get_state(state_map)
  _dJ(q) = ∫(q*...)dΩ
  Gridap.FESpaces.assemble_vector!(_dJ,dJ,V_φ)
end

This functionality is subject to change.

source

Staggered-type problems

Staggered-type problems can be handled purely with Zygote and the other existing StateMap implementations. This is preferred over the StaggeredStateMap implementations.

For example, we can solve a problem where the second FE problem depends on the first via the following:

## Weak forms
a1(u1,v1,φ) = ...
l1(v1,φ) = ...
# Treat (u1,φ) as the primal variable
a2(u2,v2,(u1,φ)) = ...
l2(v2,(u1,φ)) = ...

## Build StateMaps
φ_to_u1 = AffineFEStateMap(a1,l1,U1,V,V_φ,φh)
# u1φ_to_u2 has a MultiFieldFESpace V_u1φ of primal vars
u1φ_to_u2 = AffineFEStateMap(a2,l2,U2,V,V_u1φ,interpolate([1,φh],V_u1φ))
# The StateParamMap F needs to take a MultiFieldFEFunction u1u2h ∈ U_u1u2
F = GridapTopOpt.StateParamMap(F,U_u1u2,V_φ,assem_U_u1u2,assem_V_φ)

function φ_to_j(φ)
  u1 = φ_to_u1(φ)
  u1φ = combine_fields(V_u1φ,u1,φ) # Combine vectors of DOFs
  u2 = u1φ_to_u2(u1φ)
  u1u2 = combine_fields(U_u1u2,u1,u2)
  F(u1u2,φ)
end

pcf = CustomPDEConstrainedFunctionals(...)

GridapTopOpt + GridapEmbedded + Zygote

GridapTopOpt.CustomEmbeddedPDEConstrainedFunctionals — Type

struct CustomEmbeddedPDEConstrainedFunctionals{N,A} <: AbstractPDEConstrainedFunctionals{N}

A version of CustomPDEConstrainedFunctionals that has an embedded_collection to allow the state_map to be updated given new FE spaces for the forward problem. This is currently required for unfitted methods.

source

GridapTopOpt.CustomEmbeddedPDEConstrainedFunctionals — Method

CustomEmbeddedPDEConstrainedFunctionals(
  φ_to_jc :: Function,
  num_constraints,
  embedded_collection :: EmbeddedCollection;
  analytic_dJ = nothing,
  analytic_dC = fill(nothing,num_constraints)
)

Create an instance of CustomEmbeddedPDEConstrainedFunctionals. Here, num_constraints specifies the number of constraints.

Note

If you have one or more state_map objects in embedded_collection, you should include these as a vector under the :state_map key of the embedded_collection. This is used in the optimiser to get the current state of the maps for the uh output in

for (it, uh, φh) in optimiser
  ...
end

If you do not have a :state_map in the embedded_collection, then get_state, and the corresponding output of uh above will be nothing.

This functionality is subject to change.

source