PDE Constrained Optimisation with G-ADOPT - Field Values¶
In this tutorial, we invert for an (unknown) initial condition, from a given final state.
We will see how we can use the PDE constrained optimisation functionality of G-ADOPT to optimize one of the inputs to a PDE for a specified desired outcome. We will use a time-dependent advection-diffussion equation as our PDE and see how, for a given final state of the solution, we can retrieve what the initial condition should be, via an adjoint approach.
To turn on the adjoint, one simply imports the gadopt.inverse module to enable all taping functionality from pyadjoint. The tape automatically registers all operations that form part of the forward model, which is used to automatically form the adjoint (backward) model.
from gadopt import *
from gadopt.inverse import *
Create synthetic twin experiment with final state for a known initial condition.¶
We first run a simple advection-diffusion model with a known initial condition. Of that model we only store the solution at the final timestep, which we will use in our inversion experiment later as the target final state.
We setup the model in a form compatible with our previous examples, with a mesh, function spaces,
a prescribed velocity field, boundary conditions etc... We utilise the EnergySolver of G-ADOPT to
set up an energy equation under the Boussinesq Approximation, which is just a scalar
advection-diffusion equation for temperature.
mesh = UnitSquareMesh(40, 40)
mesh.cartesian = True
left, right, bottom, top = 1, 2, 3, 4 # Boundary IDs.
V = VectorFunctionSpace(mesh, "CG", 2) # Function space for velocity.
Q = FunctionSpace(mesh, "CG", 1) # Function space for the scalar (Temperature).
T = Function(Q, name="Temperature")
T0 = Function(Q, name="Initial_Temperature") # T Initial condition which we will invert for.
# Set up prescribed velocity field -- an anti-clockwise rotation around (0.5, 0.5):
x, y = SpatialCoordinate(mesh)
u = Function(V, name="Velocity").interpolate(as_vector((-y + 0.5, x - 0.5)))
# The Rayleigh number, Ra, is not actually used here, but we set a value for the diffusivity, kappa.
approximation = BoussinesqApproximation(Ra=1, kappa=2e-2)
delta_t = 0.1 # timestep
energy_solver = EnergySolver(T, u, approximation, delta_t, ImplicitMidpoint)
# Setup the initial condition for Temperature: a Gaussian at $(0.75, 0.5)$.
# This will be the initial condition we will try to invert for later.
x0, y0 = 0.75, 0.5
w = 0.1
r2 = (x - x0) ** 2 + (y - y0) ** 2
T0.interpolate(exp(-r2 / w**2))
Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7a3c68b72870>, FiniteElement('Lagrange', triangle, 1), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 0)), 4)
We can visualise the initial temperature field using Firedrake's built-in plotting functionality.
import matplotlib.pyplot as plt
fig, axes = plt.subplots()
collection = tripcolor(T0, axes=axes, cmap='magma', vmax=0.15)
fig.colorbar(collection);
After setting the initial condition for T, we can next run the forward model, for a specified number of timesteps. Pretty simple, huh?
num_timesteps = 10
T.project(T0)
for timestep in range(num_timesteps):
energy_solver.solve()
We can plot the final temperature solution, which you will see has been rotated whilst simultaneously diffusing.
fig, axes = plt.subplots()
collection = tripcolor(T, axes=axes, cmap='magma', vmax=0.15)
fig.colorbar(collection);
We next save this final synthetic model state in a checkpoint file for later use:
with CheckpointFile("Final_State.h5", "w") as final_checkpoint:
final_checkpoint.save_mesh(mesh)
final_checkpoint.save_function(T, name="Temperature")
Advection diffusion model with unknown initial condition¶
We now start from scratch again with an advection-diffusion model with the same configuration, except this time we don't know the initial condition. As we want to measure for different initial conditions, how well the final state of the model matches the one we just saved, we first read back that target final state. We will also use the mesh from the checkpoint file to construct the model.
with CheckpointFile("Final_State.h5", "r") as final_checkpoint:
mesh = final_checkpoint.load_mesh()
mesh.cartesian = True
T_target = final_checkpoint.load_function(mesh, name="Temperature")
With this information stored, we now set up the model exactly as before:
V = VectorFunctionSpace(mesh, "CG", 2)
Q = FunctionSpace(mesh, "CG", 1)
T = Function(Q, name="Temperature")
T0 = Function(Q, name="Initial_Temperature")
x, y = SpatialCoordinate(mesh)
u = Function(V, name="Velocity").interpolate(as_vector((-y + 0.5, x - 0.5)))
approximation = BoussinesqApproximation(Ra=1, kappa=2e-2)
delta_t = 0.1
energy_solver = EnergySolver(T, u, approximation, delta_t, ImplicitMidpoint)
# Make our solver output a little less verbose, aiding interpretation of optimisation output below:
if "ksp_converged_reason" in energy_solver.solver_parameters:
del energy_solver.solver_parameters["ksp_converged_reason"]
This time, however, we don't want to use the known initial condition. Instead we will start with
the final state from our synthetic forward model, which will then be further rotated and diffused. After
a specified number of timesteps we compute the mismatch between the predicted final state, and the
checkpointed final state from our synthetic twin. This computation, the timesteps and the mismatch calculation,
forms the forward model that we want to invert. Its adjoint will be created automatically from the tape, which
registers all operations in the model. Since the tape was automatically started at the top when we imported
gadopt.inverse, we must ensure we don't get mixed up with any operations that happened in our initial
synthetic twin model, so we first clear everything that has already been registered from the tape.
tape = get_working_tape()
tape.clear_tape()
At this stage, we are good to specify our initial guess for temperature, from T_target (i.e. the final state of our synthetic forward run). We do this by interpolating T_target to T0. Note that in theory, we could start from an arbitrary initial condition here, provided it is bounded between 0 and 1.
T0.project(T_target)
Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7a3c13d60f80>, FiniteElement('Lagrange', triangle, 1), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 55)), 126)
We want to optimise this initial temperature state, T0, and so we specify this as the control:
m = Control(T0)
Based on our guess for the initial temperature, we run the model for a specified number of timesteps:
T.project(T0)
for timestep in range(num_timesteps):
energy_solver.solve()
For good performance of optimisation algorithms, it is important to scale both the control and the functional values to be of order 1. Note that mathematically scaling the functional should not change the optimal solution. In the following lines, we define our scaled functional (J) to be the l2 misfit between predicted and true final state: