11-extract-adjoint-solutions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import firedrake\n",
    "except ImportError:\n",
    "    !wget \"https://fem-on-colab.github.io/releases/firedrake-install-real.sh\" -O \"/tmp/firedrake-install.sh\" && bash \"/tmp/firedrake-install.sh\"\n",
    "    import firedrake"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving adjoint equations using `firedrake-adjoint`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to use `firedrake-adjoint` to solve an adjoint equation and extract the solution thereof.\n",
    "\n",
    "Suppose we have a PDE in 'residual form'\n",
    "\n",
    "$$F(u)=0,\\quad u\\in V,$$\n",
    "\n",
    "for some function space $V$.\n",
    "Given an objective functional $J:V\\rightarrow\\mathbb R$, the adjoint equation is given by\n",
    "\n",
    "$$\\frac{\\partial F}{\\partial u}^T\\lambda=\\frac{\\partial J}{\\partial u}^T,\\quad\\lambda\\in V.$$\n",
    "\n",
    "We seek to compute the adjoint solution, $\\lambda$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose the PDE involves a parameter $\\nu$ which is defined as a Firedrake `Constant` or `Function` (rather than a float).\n",
    "Then `firedrake-adjoint` enables us to compute the gradient $\\mathrm dJ/\\mathrm d\\nu$.\n",
    "Whilst we might not actually need this for our application, performing the computation endows the `SolveVarFormBlock`s with adjoint solutions.\n",
    "Why is that?\n",
    "\n",
    "Taking the transpose of the adjoint equation gives\n",
    "\n",
    "$$\\frac{\\partial J}{\\partial u}=\\lambda^T\\frac{\\partial F}{\\partial u}.$$\n",
    "\n",
    "Now, expanding the gradient $\\mathrm dJ/\\mathrm d\\nu$ using the chain rule gives\n",
    "\n",
    "$$\n",
    "\\frac{\\mathrm dJ}{\\mathrm d\\nu}\n",
    "=\\frac{\\partial J}{\\partial u}\\frac{\\mathrm du}{\\mathrm d\\nu}+\\frac{\\partial J}{\\partial\\nu}\n",
    "=\\lambda^T\\frac{\\partial F}{\\partial u}\\frac{\\mathrm du}{\\mathrm d\\nu}+\\frac{\\partial J}{\\partial\\nu}\n",
    "=\\lambda^T\\frac{\\partial F}{\\partial\\nu}+\\frac{\\partial J}{\\partial\\nu}.\n",
    "$$\n",
    "\n",
    "In order to compute the gradient $\\mathrm dJ/\\mathrm d\\nu$, it is therefore sufficient to evaluate:\n",
    "1. the adjoint solution, $\\lambda$;\n",
    "2. the _partial_ derivative of the PDE residual w.r.t. the parameter $\\nu$;\n",
    "3. the _partial_ derivative of the objective functional w.r.t. $\\nu$.\n",
    "\n",
    "In order to compute the gradient, `firedrake-adjoint` does all three of these things."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The forward problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import Firedrake with adjoint mode activated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code in this cell makes plots appear an appropriate size and resolution in the browser window\n",
    "%config InlineBackend.figure_format = 'svg'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (11, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake import *\n",
    "from firedrake.adjoint import *\n",
    "continue_annotation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We employ the same finite element discretisation of Burgers equation in $\\mathbb P2$ space as in the [Firedrake demo](https://firedrakeproject.org/demos/burgers.py.html). Here $\\nu$ is the viscosity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple mesh\n",
    "n = 32\n",
    "mesh = UnitSquareMesh(n, n)\n",
    "\n",
    "# Define P2 function space and corresponding test function\n",
    "V = VectorFunctionSpace(mesh, \"Lagrange\", 2)\n",
    "v = TestFunction(V)\n",
    "\n",
    "# Create Functions for the solution and time-lagged solution\n",
    "u = Function(V, name=\"Velocity\")\n",
    "u_ = Function(V)\n",
    "\n",
    "# Assign initial condition\n",
    "x, y = SpatialCoordinate(mesh)\n",
    "u_ = interpolate(as_vector([sin(pi*x), 0]), V)\n",
    "u.assign(u_)\n",
    "\n",
    "# Set diffusivity constant\n",
    "R = FunctionSpace(mesh, \"R\", 0)\n",
    "nu = Function(R)\n",
    "nu.assign(0.0001)\n",
    "\n",
    "# Define nonlinear form\n",
    "dt = 1.0/n\n",
    "F = (inner((u - u_)/dt, v) + inner(dot(u, nabla_grad(u)), v) + nu*inner(grad(u), grad(v)))*dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having set up the residual,`F`, we are able to compute weak solutions of the PDE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Timestepping details\n",
    "end_time = 0.5\n",
    "timesteps_per_export = 4\n",
    "num_timesteps = int(end_time/dt)\n",
    "num_exports = num_timesteps//timesteps_per_export\n",
    "\n",
    "# Store forward solution at exports so we can plot again later\n",
    "forward_solutions = [u.copy(deepcopy=True)]\n",
    "\n",
    "# Time integrate\n",
    "i = 0\n",
    "t = 0.0\n",
    "while (t < end_time):\n",
    "    solve(F == 0, u)\n",
    "    u_.assign(u)\n",
    "    t += dt\n",
    "    i += 1\n",
    "    if i % timesteps_per_export == 0:\n",
    "        forward_solutions.append(u.copy(deepcopy=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot solution at each export timestep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake.pyplot import tricontourf\n",
    "\n",
    "fig, axs = plt.subplots(num_exports+1, sharex='col')\n",
    "\n",
    "for i in range(len(forward_solutions)):\n",
    "    tricontourf(forward_solutions[i], axes=axs[i])\n",
    "    axs[i].annotate('t={:.2f}'.format(i*timesteps_per_export*dt), (0.05, 0.05), color='white');\n",
    "axs[0].set_title('Forward solution, u');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The objective functional"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to solve the adjoint equation, we need an objective functional. For this example, consider the following integral over the 'outflow' (right hand) boundary at the final time $T$:\n",
    "\n",
    "$$J(u)=\\int_{\\Gamma_2}u|_{t=T}\\cdot u|_{t=T}\\;\\mathrm ds,\n",
    "\\quad\\quad\\Gamma_2=\\{(1,y)\\mid y\\in(0,1)\\}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J_form = inner(u, u)*ds(2)\n",
    "J = assemble(J_form)\n",
    "print(\"Objective value: {:.4f}\".format(J))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having solved the forward equation, `firedrake-adjoint` has all it needs to solve the adjoint equation.\n",
    "Therefore, we can tell it to stop annotating the tape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stop_annotating();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting adjoint solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the working tape and take a look at its blocks. These correspond to the (high level) operations we've performed during the simulation. For example, value assignment, FEM solves and functional assembly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tape = get_working_tape()\n",
    "for i, block in enumerate(tape._blocks):\n",
    "    print(\"Block {:2d}: {:}\".format(i, type(block)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The time integration routine is apparent, with the repeated solves (`SolveVarFormBlock`) and assignments (`FunctionAssignBlock`). Having time integrated the forward equation, we end by assembling the objective functional (`AssembleBlock`).\n",
    "\n",
    "As discussed above, we get `firedrake-adjoint` to compute the gradient $\\mathrm dJ/\\mathrm d\\nu$ in order to endow the solve blocks with adjoint solutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = compute_gradient(J, Control(nu))\n",
    "print(\"Gradient of J w.r.t. diffusivity = {:.4f}\".format(*g.dat.data_ro))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: Convince yourself that the negative gradient makes sense\n",
    "\n",
    "- Hint 1: What is it that the objective functional quantifies?\n",
    "- Hint 2: How does diffusion affect this quantity?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are only interested in the blocks corresponding to nonlinear solves.\n",
    "These may be extracted using Firedrake's helper function `get_solve_blocks`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake.adjoint import get_solve_blocks\n",
    "\n",
    "solve_blocks = get_solve_blocks()\n",
    "assert len(solve_blocks) == num_timesteps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each solve block, `block`, the adjoint solution is stored as the attribute `block.adj_sol`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for block in solve_blocks:\n",
    "    assert block.adj_sol is not None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By accessing these, we are now able to plot the adjoint solution.\n",
    "Remember that information propagates _backwards_ in time for the adjoint, so we should view the plots in ascending order.\n",
    "\n",
    "Since the objective functional is defined at the end time, its derivative provides an 'initial condition' for the adjoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 'Initial condition' for both adjoint\n",
    "dJdu = assemble(derivative(J_form, u))\n",
    "\n",
    "# Plot adjoint solutions at matching timesteps to forward\n",
    "fig, axs = plt.subplots(num_exports+1, 2, sharex='col', figsize=(10, 5))\n",
    "for i in range(num_exports+1):\n",
    "    t = i*timesteps_per_export*dt\n",
    "    tricontourf(forward_solutions[i], axes=axs[i, 0])\n",
    "    adjoint_solution = dJdu if i == num_exports else solve_blocks[timesteps_per_export*i].adj_sol\n",
    "    # Get the Riesz representer\n",
    "    adjoint_solution = dJdu.riesz_representation(riesz_map=\"H1\")\n",
    "    tricontourf(adjoint_solution, axes=axs[i, 1])\n",
    "    axs[i, 0].annotate('t={:.2f}'.format(t), (0.05, 0.05), color='white');\n",
    "    axs[i, 1].annotate('t={:.2f}'.format(t), (0.05, 0.05), color='white');\n",
    "\n",
    "# Plot formatting\n",
    "axs[0, 0].set_title('Forward solution');\n",
    "axs[0, 1].set_title('Adjoint solution');"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}