{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "258cd78c-55e3-40e6-97cb-057064281a61",
   "metadata": {},
   "source": [
    "# Functional minimization\n",
    "\n",
    "Gradient-based optimization algorithms attempt to use derivative information, for example as computed using the adjoint method, to accelerate the solution of optimization problems. A typical example is to seek the minimum of an objective cost functional, constrained using one or more partial differential equations. This notebook describes how to solve such partial differential equation constrained optimization problems using tlm_adjoint.\n",
    "\n",
    "This example makes use of tlm_adjoint with the [Firedrake](https://firedrakeproject.org/) backend, and we assume real spaces and a real build of Firedrake throughout. The minimization problem is solved using the [Toolkit for Advanced Optimization (TAO)](https://petsc.org/release/manual/tao/).\n",
    "\n",
    "## Forward problem\n",
    "\n",
    "We first construct a partial differential equation constrained optimization problem where we know the answer – where we know the value of the control which minimizes the objective functional subject to the constraint that the forward problem is solved. To do this we consider the solution of the modified Helmholtz equation in the unit square domain,\n",
    "\n",
    "$$\\alpha^2 \\nabla^2 u - u = m \\qquad \\text{on} ~ \\Omega = \\left( 0, 1 \\right)^2,$$\n",
    "\n",
    "where $\\alpha \\in \\mathbb{R}$, subject to doubly periodic boundary conditions. We define an objective cost functional\n",
    "\n",
    "$$J \\left( u, m \\right) = \\frac{1}{2} \\int_\\Omega \\left( u - \\tilde{u} \\right)^2 + \\frac{1}{2} \\beta^2 \\int_\\Omega \\left( m - \\tilde{m} \\right)^2,$$\n",
    "\n",
    "where $\\tilde{u}$ and $\\tilde{m}$ are given functions and $\\beta \\ne 0$ is a real scalar. If we set\n",
    "\n",
    "$$\\tilde{m} \\left( x, y \\right) = -\\sin \\left( 2 \\pi x \\right) \\sin \\left( 2 \\pi y \\right),$$\n",
    "$$\\tilde{u} \\left( x, y \\right) = -\\frac{1}{1 + 8 \\pi^2 \\alpha^2} m \\left( x, y \\right),$$\n",
    "\n",
    "where $x$ and $y$ are Cartesian coordinates in $\\mathbb{R}^2$, then $u = \\tilde{u}$ and $m = \\tilde{m}$ will be the minimum of $J$ where the modified Helmholtz problem is solved.\n",
    "\n",
    "We consider a continuous Galerkin finite element discretization, seeking $u \\in V$ such that\n",
    "\n",
    "$$\\forall \\zeta \\in V \\qquad \\alpha^2 \\int_\\Omega \\nabla \\zeta \\cdot \\nabla u + \\int_\\Omega \\zeta u = -\\int_\\Omega \\zeta m,$$\n",
    "\n",
    "where $V$ is a real continuous $P_1$ finite element space whose elements satisfy the doubly periodic boundary conditions. We now define $\\tilde{m} \\in V$ and $\\tilde{u} \\in V$ via interpolation, at mesh vertices, of the functions given above.\n",
    "\n",
    "We first use Firedrake to solve the forward problem. We consider $\\alpha = 0.1$, $\\beta = 0.1$, and an 'initial guess' where $m = -1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f41c40c7-d0c3-438e-b874-e35cfeee7d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from firedrake import *\n",
    "from tlm_adjoint.firedrake import *\n",
    "from firedrake.pyplot import tricontourf\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "alpha = Constant(0.1)\n",
    "beta = Constant(0.1)\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    space = m.function_space()\n",
    "    X = SpatialCoordinate(space.mesh())\n",
    "    test, trial = TestFunction(space), TrialFunction(space)\n",
    "\n",
    "    m_tilde = Function(space, name=\"m_tilde\").interpolate(\n",
    "        -sin(2 * pi * X[0]) * sin(2 * pi * X[1]))\n",
    "    u_tilde = Function(space, name=\"u_tilde\").assign(\n",
    "        -(1.0 / (1.0 + 8.0 * pi * pi * alpha * alpha)) * m_tilde)\n",
    "\n",
    "    u = Function(space, name=\"u\")\n",
    "    solve(alpha * alpha * inner(grad(trial), grad(test)) * dx + inner(trial, test) * dx\n",
    "          == -inner(m, test) * dx,\n",
    "          u, solver_parameters={\"ksp_type\": \"cg\",\n",
    "                                \"pc_type\": \"sor\",\n",
    "                                \"ksp_atol\": 1.0e-32,\n",
    "                                \"ksp_rtol\": 1.0e-12})\n",
    "\n",
    "    J = Functional(name=\"J\")\n",
    "    J.addto(0.5 * inner(u - u_tilde, u - u_tilde) * dx)\n",
    "    J.addto(0.5 * beta * beta * inner(m - m_tilde, m - m_tilde) * dx)\n",
    "    return m_tilde, u, J\n",
    "\n",
    "\n",
    "mesh = PeriodicSquareMesh(20, 20, 1.0)\n",
    "X = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "\n",
    "m_0 = Function(space, name=\"m_0\").interpolate(Constant(-1.0))\n",
    "m_tilde, u, J = forward(m_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b81c75-ceba-4e7e-a5ed-c3cbb521617c",
   "metadata": {},
   "source": [
    "## Inverse problem\n",
    "\n",
    "We now seek to solve the inverse problem: to find the $m \\in V$ which minimizes $J$, subject to the discretized modified Helmholtz problem being solved.\n",
    "\n",
    "In the following we use the Toolkit for Advanced Optimization (TAO) to solve the inverse problem. We use the Limited-Memory Variable Metric (LMVM) approach with an absolute tolerance of $10^{-10}$. Noting that the adjoint computed derivative is an element of the dual space $V^*$, we need to define an appropriate dual space inner product. Here we define an inner product using the inverse mass matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d2768b-5ff1-4f58-9370-1833e01cb612",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_J(m):\n",
    "    _, _, J = forward(m)\n",
    "    return J\n",
    "\n",
    "\n",
    "M_solver = LinearSolver(assemble(inner(TrialFunction(space), TestFunction(space)) * dx),\n",
    "                        solver_parameters={\"ksp_type\": \"cg\",\n",
    "                                           \"pc_type\": \"sor\",\n",
    "                                           \"ksp_atol\": 1.0e-32,\n",
    "                                           \"ksp_rtol\": 1.0e-12})\n",
    "\n",
    "\n",
    "def M_inv_action(x):\n",
    "    y = Function(space)\n",
    "    M_solver.solve(y, x.copy(deepcopy=True))\n",
    "    return y\n",
    "\n",
    "\n",
    "def post_callback(tao):\n",
    "    print(f\"{tao.getIterationNumber()=}\")\n",
    "\n",
    "\n",
    "m = minimize_tao(forward_J, m_0, gatol=1.0e-10, grtol=0.0,\n",
    "                 M_inv_action=M_inv_action, post_callback=post_callback)\n",
    "m.rename(name=\"m\")\n",
    "m_tilde, u, J = forward(m)\n",
    "\n",
    "\n",
    "def plot_output(u, title):\n",
    "    r = (u.dat.data_ro.min(), u.dat.data_ro.max())\n",
    "    eps = (r[1] - r[0]) * 1.0e-12\n",
    "    p = tricontourf(u, np.linspace(r[0] - eps, r[1] + eps, 32))\n",
    "    plt.gca().set_title(title)\n",
    "    plt.colorbar(p)\n",
    "    plt.gca().set_aspect(1.0)\n",
    "\n",
    "\n",
    "plot_output(u, title=\"u\")\n",
    "plot_output(m, title=\"m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2bf05e9-0e6e-4a4c-8b93-9e6019f689ee",
   "metadata": {},
   "source": [
    "We now test the inverse procedure by checking that it converges to the expected result, considering meshes with decreasing element size. We compute the $L^2$ error in each case, and estimate the order of convergence by a power law fit between the error norms computed using subsequent pairs of meshes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebfbaff7-a877-4e9b-bac7-b7bdad621e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ns = np.array([20 * (2 ** p) for p in range(4)], dtype=int)\n",
    "error_norms = []\n",
    "for N in Ns:\n",
    "    mesh = PeriodicSquareMesh(N, N, 1.0)\n",
    "    X = SpatialCoordinate(mesh)\n",
    "    space = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "\n",
    "    m_0 = Function(space, name=\"m_0\").interpolate(Constant(-1.0))\n",
    "\n",
    "    M_solver = LinearSolver(assemble(inner(TrialFunction(space), TestFunction(space)) * dx),\n",
    "                            solver_parameters={\"ksp_type\": \"cg\",\n",
    "                                               \"pc_type\": \"sor\",\n",
    "                                               \"ksp_atol\": 1.0e-32,\n",
    "                                               \"ksp_rtol\": 1.0e-12})\n",
    "\n",
    "    def M_inv_action(x):\n",
    "        y = Function(space)\n",
    "        M_solver.solve(y, x.copy(deepcopy=True))\n",
    "        return y\n",
    "\n",
    "    m = minimize_tao(forward_J, m_0, gatol=1.0e-10, grtol=0.0,\n",
    "                     M_inv_action=M_inv_action, post_callback=post_callback)\n",
    "    m_tilde, u, J = forward(m)\n",
    "\n",
    "    m_error_norm = sqrt(abs(assemble(inner(m - m_tilde, m - m_tilde) * dx)))\n",
    "    print(f\"{N=} {m_error_norm=}\")\n",
    "    error_norms.append(m_error_norm)\n",
    "error_norms = np.array(error_norms, dtype=float)\n",
    "\n",
    "orders = -np.log(error_norms[1:] / error_norms[:-1]) / np.log(Ns[1:] / Ns[:-1])\n",
    "print(f\"{orders=}\")\n",
    "\n",
    "assert (orders > 1.99).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a913c18-c76b-4119-aab5-db6176ddac04",
   "metadata": {},
   "source": [
    "We find that we have close to second order convergence."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}