{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "82be80f5-acad-4947-8f9a-c6a7a87a7ace",
   "metadata": {},
   "source": [
    "# Defining custom operations\n",
    "\n",
    "This notebook describes the tlm_adjoint 'escape hatch' mechanism, which allows custom operations to be defined, recorded by the internal manager, and used in higher order derivative calculations.\n",
    "\n",
    "The approach used by tlm_adjoint to define custom operations is described in\n",
    "\n",
    "- James R. Maddison, Daniel N. Goldberg, and Benjamin D. Goddard, 'Automated calculation of higher order partial differential equation constrained derivative information', SIAM Journal on Scientific Computing, 41(5), pp. C417–C445, 2019, doi: 10.1137/18M1209465\n",
    "\n",
    "We assume real spaces and a real build of Firedrake throughout.\n",
    "\n",
    "## Forward problem\n",
    "\n",
    "We consider the Poisson equation in the unit square domain\n",
    "\n",
    "$$\\nabla^2 u = m \\quad \\text{on} ~ \\Omega = \\left( 0, 1 \\right)^2,$$\n",
    "\n",
    "with\n",
    "\n",
    "$$\\int_\\Omega m = 0,$$\n",
    "\n",
    "subject to boundary conditions\n",
    "\n",
    "$$\\nabla u \\cdot \\hat{n} = 0,$$\n",
    "\n",
    "where $\\hat{n}$ is an outward unit normal on the boundary $\\partial \\Omega$ of the domain $\\Omega$.\n",
    "\n",
    "We consider a continuous Galerkin discretization, now seeking $u \\in V$ such that\n",
    "\n",
    "$$\\forall \\zeta \\in V \\qquad \\int_\\Omega \\nabla \\zeta \\cdot \\nabla u = -\\int_\\Omega \\zeta m,$$\n",
    "\n",
    "where $V$ is a real $P_1$ continuous finite element space defining functions on the domain $\\Omega = \\left( 0, 1 \\right)^2$. $m \\in V$ is now defined via\n",
    "\n",
    "$$m = \\mathcal{I} \\left[ \\cos \\left( \\pi x \\right) \\cos \\left( \\pi y \\right) \\right] - k,$$\n",
    "\n",
    "where $\\mathcal{I}$ maps to an element of $V$ through interpolation at mesh vertices, and where $k$ is defined so that $\\int_\\Omega m = 0$.\n",
    "\n",
    "As stated the discrete problem does not have a unique solution – given any solution $u_0$ to the problem we can define a new solution $u_1 = u_0 + c_1$ for any scalar $c_1$. We need to identify a solution by supplying an appropriate constraint. We will use the constraint\n",
    "\n",
    "$$\\sum_i \\tilde{u}_i = 0,$$\n",
    "\n",
    "where the $\\tilde{u}_i$ are the degrees of freedom associated with $u \\in V$.\n",
    "\n",
    "An implementation using Firedrake, solving for $u$ and then computing $\\frac{1}{2} \\int_\\Omega u^2$, takes the form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17707b1c-9237-46a3-8c70-f1e700c8ad00",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from firedrake import *\n",
    "from firedrake.pyplot import tricontourf\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "mesh = UnitSquareMesh(10, 10)\n",
    "X = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "test, trial = TestFunction(space), TrialFunction(space)\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    k = Constant(assemble(m * dx) / assemble(Constant(1.0) * dx(mesh)))\n",
    "    m_tilde = Function(m.function_space(), name=\"m_tilde\").interpolate(m - k)\n",
    "    \n",
    "    u = Function(space, name=\"u\")\n",
    "    nullspace = VectorSpaceBasis(comm=u.comm, constant=True)\n",
    "\n",
    "    solve(inner(grad(trial), grad(test)) * dx == -inner(m_tilde, test) * dx,\n",
    "          u, nullspace=nullspace, transpose_nullspace=nullspace,\n",
    "          solver_parameters={\"ksp_type\": \"cg\", \"pc_type\": \"hypre\",\n",
    "                             \"pc_hypre_type\": \"boomeramg\",\n",
    "                             \"ksp_rtol\": 1.0e-10, \"ksp_atol\": 1.0e-16})\n",
    "\n",
    "    J = assemble(0.5 * inner(u, u) * dx)\n",
    "    return u, J\n",
    "\n",
    "\n",
    "m = Function(space, name=\"m\").interpolate(cos(pi * X[0]) * cos(pi * X[1]))\n",
    "\n",
    "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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc5e0b3d-426c-4b75-8b74-d46037c5e822",
   "metadata": {},
   "source": [
    "## Adding tlm_adjoint\n",
    "\n",
    "Next we add the use of tlm_adjoint, which requires some minor modification of the forward code.\n",
    "\n",
    "`compute_gradient` is used to compute the derivative of $\\frac{1}{2} \\int_\\Omega u^2$ with respect to $m$, subject to the contraint that the forward problem is solved. The derivative is visualized using a Riesz map defined by the $L^2$ inner product. A Taylor remainder convergence test is used to verify the derivative.\n",
    "\n",
    "Note that, given an arbitrary $m \\in V$, we define\n",
    "\n",
    "$$\\tilde{m} = m - \\frac{\\int_\\Omega m}{\\int_\\Omega 1},$$\n",
    "\n",
    "so that $\\int_\\Omega \\tilde{m} = 0$, and use $\\tilde{m}$ in place of $m$ when solving the discrete Poisson problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa48c271-64ab-4104-8d8d-2ebc46adb7f0",
   "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",
    "np.random.seed(54151610)\n",
    "reset_manager()\n",
    "\n",
    "mesh = UnitSquareMesh(10, 10)\n",
    "X = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "test, trial = TestFunction(space), TrialFunction(space)\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    v = Constant(name=\"v\")\n",
    "    Assembly(v, Constant(1.0) * dx(mesh)).solve()\n",
    "    k = Constant(name=\"k\")\n",
    "    Assembly(k, (m / v) * dx).solve()\n",
    "    m_tilde = Function(m.function_space(), name=\"m_tilde\").interpolate(m - k)\n",
    "    \n",
    "    u = Function(space, name=\"u\")\n",
    "    nullspace = VectorSpaceBasis(comm=u.comm, constant=True)\n",
    "\n",
    "    solve(inner(grad(trial), grad(test)) * dx == -inner(m_tilde, test) * dx,\n",
    "          u, nullspace=nullspace, transpose_nullspace=nullspace,\n",
    "          solver_parameters={\"ksp_type\": \"cg\", \"pc_type\": \"hypre\",\n",
    "                             \"pc_hypre_type\": \"boomeramg\",\n",
    "                             \"ksp_rtol\": 1.0e-10, \"ksp_atol\": 1.0e-16})\n",
    "\n",
    "    J = Functional(name=\"J\")\n",
    "    J.assign(0.5 * inner(u, u) * dx)\n",
    "    return u, J\n",
    "\n",
    "\n",
    "m = Function(space, name=\"m\").interpolate(cos(pi * X[0]) * cos(pi * X[1]))\n",
    "\n",
    "start_manager()\n",
    "u, J = forward(m)\n",
    "stop_manager()\n",
    "\n",
    "dJ = compute_gradient(J, 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(dJ.riesz_representation(\"L2\"), title=r\"$g^\\sharp$\")\n",
    "\n",
    "\n",
    "def forward_J(m):\n",
    "    _, J = forward(m)\n",
    "    return J\n",
    "\n",
    "\n",
    "min_order = taylor_test(forward_J, m, J_val=J.value, dJ=dJ)\n",
    "assert min_order > 1.99"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23af6882-f085-4fa5-ae75-3ff04b367e3c",
   "metadata": {},
   "source": [
    "## Defining a custom `Equation`\n",
    "\n",
    "Since the $u$ we have computed should satisfy\n",
    "\n",
    "$$\\sum_i \\tilde{u}_i = 0,$$\n",
    "\n",
    "we should have that the derivative of this quantity with respect to $m$, subject to the forward problem being solved, is zero. We will compute this derivative by defining a custom operation which sums the degrees of freedom of a finite element discretized function. We can do this by inheriting from the `Equation` class and implementing appropriate methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3bdc61-be2a-4e94-82f7-d8f6d5faa92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tlm_adjoint import Equation\n",
    "\n",
    "\n",
    "class Sum(Equation):\n",
    "    def __init__(self, x, y):\n",
    "        super().__init__(x, deps=(x, y), nl_deps=(), ic=False, adj_ic=False,\n",
    "                         adj_type=\"conjugate_dual\")\n",
    "\n",
    "    def forward_solve(self, x, deps=None):\n",
    "        if deps is None:\n",
    "            deps = self.dependencies()\n",
    "        _, y = deps\n",
    "        with y.dat.vec_ro as y_v:\n",
    "            y_sum = y_v.sum()\n",
    "        x.assign(y_sum)\n",
    "\n",
    "    def adjoint_jacobian_solve(self, adj_x, nl_deps, b):\n",
    "       return b\n",
    "\n",
    "    def adjoint_derivative_action(self, nl_deps, dep_index, adj_x):\n",
    "        if dep_index != 1:\n",
    "            raise ValueError(\"Unexpected dep_index\")\n",
    "\n",
    "        _, y = self.dependencies()\n",
    "        b = Cofunction(var_space(y).dual())\n",
    "        adj_x = float(adj_x)\n",
    "        b.dat.data[:] = -adj_x\n",
    "        return b\n",
    "\n",
    "    def tangent_linear(self, tlm_map):\n",
    "        tau_x, tau_y = (tlm_map[dep] for dep in self.dependencies())\n",
    "        if tau_y is None:\n",
    "            return ZeroAssignment(tau_x)\n",
    "        else:\n",
    "            return Sum(tau_x, tau_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "854d0055-7138-4728-a775-38cc8fecabf5",
   "metadata": {},
   "source": [
    "We'll investgate how this is put together shortly.\n",
    "\n",
    "First, we check that the `Sum` class works. In the following we check that we can compute the sum of the degrees of freedom of a function, and then use Taylor remainder convergence tests to verify a first order tangent-linear calculation, a first order adjoint calculation, and a reverse-over-forward adjoint calculation of a Hessian action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed98c56-fa77-499c-8243-8f83727c2aaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake import *\n",
    "from tlm_adjoint.firedrake import *\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(13561700)\n",
    "reset_manager()\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    m_sum = Float(name=\"m_sum\")\n",
    "    Sum(m_sum, m).solve()\n",
    "\n",
    "    J = m_sum ** 3\n",
    "    \n",
    "    return m_sum, J\n",
    "\n",
    "\n",
    "def forward_J(m):\n",
    "    _, J = forward(m)\n",
    "    return J\n",
    "\n",
    "\n",
    "mesh = UnitIntervalMesh(10)\n",
    "x, = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Discontinuous Lagrange\", 0)\n",
    "\n",
    "m = Function(space, name=\"m\")\n",
    "m.interpolate(exp(-x) * sin(pi * x))\n",
    "with m.dat.vec_ro as m_v:\n",
    "    m_sum_ref = m_v.sum()\n",
    "\n",
    "start_manager()\n",
    "m_sum, J = forward(m)\n",
    "stop_manager()\n",
    "\n",
    "# Verify the forward calculation\n",
    "\n",
    "print(f\"{float(m_sum)=}\")\n",
    "print(f\"{m_sum_ref=}\")\n",
    "assert abs(float(m_sum) - m_sum_ref) == 0.0\n",
    "\n",
    "# Verify tangent-linear and adjoint calculations\n",
    "\n",
    "min_order = taylor_test_tlm(forward_J, m, tlm_order=1)\n",
    "assert min_order > 1.99\n",
    "\n",
    "min_order = taylor_test_tlm_adjoint(forward_J, m, adjoint_order=1)\n",
    "assert min_order > 1.99\n",
    "\n",
    "min_order = taylor_test_tlm_adjoint(forward_J, m, adjoint_order=2)\n",
    "assert min_order > 1.99"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6236cab-0eb6-4fad-ab87-e71e5ab071a2",
   "metadata": {},
   "source": [
    "We can now use `Sum` with the Poisson solver. We find, as expected, that the derivative of the sum of the degrees of freedom for $u$ with respect to $m$, subject to the forward problem being solved, is zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa84206-fdc1-4672-83df-190fc974fa43",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from firedrake import *\n",
    "from tlm_adjoint.firedrake import *\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "reset_manager()\n",
    "\n",
    "mesh = UnitSquareMesh(10, 10)\n",
    "X = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "test, trial = TestFunction(space), TrialFunction(space)\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    v = Constant(name=\"v\")\n",
    "    Assembly(v, Constant(1.0) * dx(mesh)).solve()\n",
    "    k = Constant(name=\"k\")\n",
    "    Assembly(k, (m / v) * dx).solve()\n",
    "    m_tilde = Function(m.function_space(), name=\"m_tilde\").interpolate(m - k)\n",
    "    \n",
    "    u = Function(space, name=\"u\")\n",
    "    nullspace = VectorSpaceBasis(comm=u.comm, constant=True)\n",
    "\n",
    "    solve(inner(grad(trial), grad(test)) * dx == -inner(m_tilde, test) * dx,\n",
    "          u, nullspace=nullspace, transpose_nullspace=nullspace,\n",
    "          solver_parameters={\"ksp_type\": \"cg\", \"pc_type\": \"hypre\",\n",
    "                             \"pc_hypre_type\": \"boomeramg\",\n",
    "                             \"ksp_rtol\": 1.0e-10, \"ksp_atol\": 1.0e-16})\n",
    "\n",
    "    J = Functional(name=\"J\")\n",
    "    J.assign(0.5 * inner(u, u) * dx)\n",
    "    \n",
    "    u_sum = Constant(name=\"u_sum\")\n",
    "    Sum(u_sum, u).solve()\n",
    "\n",
    "    K = Functional(name=\"K\")\n",
    "    K.assign(u_sum)\n",
    "\n",
    "    return u, J, K\n",
    "\n",
    "\n",
    "m = Function(space, name=\"m\").interpolate(cos(pi * X[0]) * cos(pi * X[1]))\n",
    "\n",
    "start_manager()\n",
    "u, J, K = forward(m)\n",
    "stop_manager()\n",
    "\n",
    "dJ, dK = compute_gradient((J, K), m)\n",
    "\n",
    "dK_norm = abs(dK.dat.data_ro).max()\n",
    "print(f\"{dK_norm=}\")\n",
    "\n",
    "assert dK_norm == 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7929c30-35c5-4ac7-bf43-3761685fd03b",
   "metadata": {},
   "source": [
    "## `Equation` methods\n",
    "\n",
    "We now return to the definition of the `Sum` class, and consider each method in turn.\n",
    "\n",
    "### `__init__`\n",
    "\n",
    "```\n",
    "def __init__(self, x, y):\n",
    "    super().__init__(x, deps=(x, y), nl_deps=(), ic=False, adj_ic=False,\n",
    "                     adj_type=\"conjugate_dual\")\n",
    "```\n",
    "\n",
    "The arguments passed to the base class constructor are:\n",
    "\n",
    "- The output of the forward operation, `x`.\n",
    "- `deps`: Defines all inputs and outputs of the forward operation.\n",
    "- `nl_deps`: Defines elements of `deps` on which the associated adjoint calculations can depend.\n",
    "- `ic`: Whether the forward operation accepts a non-zero 'initial guess'.\n",
    "- `adj_ic`: Whether an adjoint solve accepts a non-zero 'initial guess'.\n",
    "- `adj_type`: Either `\"primal\"` or `\"conjugate_dual\"` defining the space for an associated adjoint variable.\n",
    "\n",
    "For this operation there is one output `x` and one input `y`. The operation is defined by the forward residual function\n",
    "\n",
    "$$F \\left( x, y \\right) = x - \\sum_i \\tilde{u}_i.$$\n",
    "\n",
    "Given a value for $y$ the output of the operation is defined to be the $x$ for which $F \\left( x, y \\right) = 0$. $F$ depends linearly on both $x$ and $y$, and so associated adjoint calculations are independent of both $x$ and $y$. Hence we set `deps=(x, y)` and `nl_deps=()`.\n",
    "\n",
    "`ic` and `adj_ic` are used to indicate whether non-zero initial guesses should be supplied to forward and adjoint iterative solvers respectively. Here we do not use iterative solvers, and so no non-zero initial guesses are needed. Hence we set `ic=False` and `adj_ic=False`.\n",
    "\n",
    "`adj_type` indicates whether an adjoint variable associated with the operation is:\n",
    "\n",
    "- In the same space as `x`. This is the case where $F$ maps to an element in the antidual space associated with $x$, and is indicated with `adj_x_type=\"primal\"`. A typical example of this case is an operation defined via the solution of a finite element variational problem.\n",
    "- In the antidual space associated with the space for `x`. This is the case where $F$ maps to an element in same space as $x$, and is indiciated with `adj_x_type=\"conjugate_dual\"`.\n",
    "\n",
    "### `forward_solve`\n",
    "\n",
    "```\n",
    "def forward_solve(self, x, deps=None):\n",
    "    if deps is None:\n",
    "        deps = self.dependencies()\n",
    "    _, y = deps\n",
    "    with y.dat.vec_ro as y_v:\n",
    "        y_sum = y_v.sum()\n",
    "    x.assign(y_sum)\n",
    "```\n",
    "\n",
    "This method computes the output of the forward operation, storing the result in `x`. `deps` is used in rerunning the forward calculation when making use of a checkpointing schedule. If supplied, `deps` contains values associated with the dependencies as defined in the constructor (defined by the `deps` argument passed to the base class constructor). Otherwise the values contained in `self.dependencies()` are used.\n",
    "\n",
    "### `adjoint_jacobian_solve`\n",
    "\n",
    "```\n",
    "def adjoint_jacobian_solve(self, adj_x, nl_deps, b):\n",
    "   return b\n",
    "```\n",
    "\n",
    "Solves a linear problem\n",
    "\n",
    "$$\\frac{\\partial F}{\\partial x}^T \\lambda = b,$$\n",
    "\n",
    "for the adjoint solution $\\lambda$, returning the result. The right-hand-side $b$ is defined by the argument `b`. In this example $\\partial F / \\partial x$ is simply the identity, and so $\\lambda = b$.\n",
    "\n",
    "`adj_x` can contain an initial guess for an iterative solver used to compute $\\lambda$. `nl_deps` contains values associated with the forward dependencies of the adjoint, as defined in the constructor (defined by the `nl_deps` argument passed to the base class constructor).\n",
    "\n",
    "### `adjoint_derivative_action`\n",
    "\n",
    "```\n",
    "def adjoint_derivative_action(self, nl_deps, dep_index, adj_x):\n",
    "    if dep_index != 1:\n",
    "        raise ValueError(\"Unexpected dep_index\")\n",
    "\n",
    "    _, y = self.dependencies()\n",
    "    b = Cofunction(var_space(y).dual())\n",
    "    adj_x = float(adj_x)\n",
    "    b.dat.data[:] = -adj_x\n",
    "    return b\n",
    "```\n",
    "\n",
    "Computes\n",
    "\n",
    "$$\\frac{\\partial F}{\\partial v}^T \\lambda,$$\n",
    "\n",
    "where $v$ is the `dep_index`th dependency (for the dependencies defined by the `deps` argument passed to the base class constructor). Again `nl_deps` contains values associated with the forward dependencies of the adjoint, as defined in the constructor.\n",
    "\n",
    "### `tangent_linear`\n",
    "\n",
    "```\n",
    "def tangent_linear(self, tlm_map):\n",
    "    tau_x, tau_y = (tlm_map[dep] for dep in self.dependencies())\n",
    "    if tau_y is None:\n",
    "        return ZeroAssignment(tau_x)\n",
    "    else:\n",
    "        return Sum(tau_x, tau_y)\n",
    "```\n",
    "\n",
    "Constructs a tangent-linear operation, for a tangent-linear computing directional derivatives with respect to the control defined by `M` with direction defined by `dM`. `tlm_map` is used to access tangent-linear variables.\n",
    "\n",
    "The new operation is defined by the residual function\n",
    "\n",
    "$$F_\\tau \\left( \\tau_x, \\tau_y \\right) = \\frac{\\partial F}{\\partial x} \\tau_x + \\frac{\\partial F}{\\partial y} \\tau_y,$$\n",
    "\n",
    "where $\\tau_x$ and $\\tau_y$ are the tangent-linear variables associated with $x$ and $y$ respectively. Given a value for $\\tau_y$ the output of the operation is defined to be the $\\tau_x$ for which $F_\\tau \\left( \\tau_x, \\tau_y \\right) = 0$.\n",
    "\n",
    "Note that this method returns an `Equation` instance – here either a `ZeroAssignment`, for the case where the tangent-linear operation sets a tangent-linear variable equal to zero, or another `Sum`. This new operation can then be recorded by the internal manager, so that reverse-over-forward algorithmic differentiation can be applied.\n",
    "\n",
    "### Reference dropping\n",
    "\n",
    "An important method, not defined in this example, is the `drop_references` method. This method is used to drop references to forward data. For example if an `Equation` subclass defines a `Function` attribute `_function`, then the `drop_references` method might look like\n",
    "\n",
    "```\n",
    "def drop_references(self):\n",
    "    super().drop_references()\n",
    "    self._function = var_replacement(self._function)\n",
    "```\n",
    "\n",
    "Here `var_replacement` returns a 'symbolic only' version of `self._function`, which can for example be used in UFL expressions, but which has no associated value.\n",
    "\n",
    "## Defining custom operations using JAX\n",
    "\n",
    "In some cases it may be more convenient to directly implement an operation using lower level code. tlm_adjoint integrates with [JAX](https://jax.readthedocs.io) to allow this. For example a Firedrake `Function` can be converted to a JAX array using\n",
    "\n",
    "```\n",
    "v = to_jax(x)\n",
    "```\n",
    "\n",
    "and the result converted back to a Firedrake `Function` using\n",
    "\n",
    "```\n",
    "x = to_firedrake(v, space)\n",
    "```\n",
    "\n",
    "Here we use JAX to compute the sum of the degrees of freedom for a Firedrake function (assuming a serial calculation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98042cd-c9b5-4ac8-abd5-8fbb1a5717ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake import *\n",
    "from tlm_adjoint.firedrake import *\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(81976463)\n",
    "reset_manager()\n",
    "\n",
    "\n",
    "def forward(m):\n",
    "    m_sum = new_jax_float(name=\"m_sum\")\n",
    "    call_jax(m_sum, to_jax(m), lambda v: v.sum())\n",
    "\n",
    "    J = m_sum ** 3\n",
    "    \n",
    "    return m_sum, J\n",
    "\n",
    "\n",
    "def forward_J(m):\n",
    "    _, J = forward(m)\n",
    "    return J\n",
    "\n",
    "\n",
    "mesh = UnitIntervalMesh(10)\n",
    "x, = SpatialCoordinate(mesh)\n",
    "space = FunctionSpace(mesh, \"Discontinuous Lagrange\", 0)\n",
    "\n",
    "m = Function(space, name=\"m\")\n",
    "m.interpolate(exp(-x) * sin(pi * x))\n",
    "with m.dat.vec_ro as m_v:\n",
    "    m_sum_ref = m_v.sum()\n",
    "\n",
    "start_manager()\n",
    "m_sum, J = forward(m)\n",
    "stop_manager()\n",
    "\n",
    "# Verify the forward calculation\n",
    "\n",
    "print(f\"{float(m_sum)=}\")\n",
    "print(f\"{m_sum_ref=}\")\n",
    "assert abs(float(m_sum) - m_sum_ref) < 1.0e-15\n",
    "\n",
    "# Verify tangent-linear and adjoint calculations\n",
    "\n",
    "min_order = taylor_test_tlm(forward_J, m, tlm_order=1)\n",
    "assert min_order > 1.99\n",
    "\n",
    "min_order = taylor_test_tlm_adjoint(forward_J, m, adjoint_order=1)\n",
    "assert min_order > 1.99\n",
    "\n",
    "min_order = taylor_test_tlm_adjoint(forward_J, m, adjoint_order=2)\n",
    "assert min_order > 1.99"
   ]
  }
 ],
 "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
}