09-hybridisation.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-release-real.sh\" -O \"/tmp/firedrake-install.sh\" && bash \"/tmp/firedrake-install.sh\"\n",
    "    import firedrake"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Static condensation and hybridization\n",
    "\n",
    "In this notebook, we explore the more advanced capabilities of Firedrake which facilitate the global reduction of finite element systems. The procedure, known as element-wise \"static condensation,\" is well-known within the FEM community. This approach requires the algebraic manipulation of locally assembled matrices and vectors. We will apply this within the context of hybridizing a mixed method.\n",
    "\n",
    "As our running example, we consider the following saddle-point system: find $(\\mathbf{u}, D) \\in V \\times U \\subset H(\\text{div}) \\times L^2$ such that\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\color{#800020}{\\int_\\Omega \\mathbf{w}\\cdot\\mathbf{u}\\,\\mathrm{d}x +\n",
    "    \\beta\\int_\\Omega \\mathbf{w}\\cdot f\\mathbf{u}^\\perp\\,\\mathrm{d}x}\n",
    "    - \\color{#2A52BE}{g\\beta\\int_\\Omega D\\nabla\\cdot\\mathbf{w}\\,\\mathrm{d}x} &= -R\\lbrack \\mathbf{w} \\rbrack,\n",
    "    \\quad \\forall \\mathbf{w} \\in V, \\\\\n",
    "    \\color{#2A52BE}{H\\beta\\int_\\Omega \\phi\\nabla\\cdot\\mathbf{u}\\,\\mathrm{d}x}\n",
    "    + \\color{#CC5500}{\\int_\\Omega \\phi D\\,\\mathrm{d}x} &= -R\\lbrack \\phi \\rbrack,\n",
    "    \\quad \\forall \\phi \\in U,\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where the residual co-vectors are defined as:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    R\\lbrack \\mathbf{w} \\rbrack &= \\int_\\Omega \\mathbf{w}\\cdot\\mathbf{u}^{0}\\,\\mathrm{d}x -\n",
    "    \\beta\\int_\\Omega \\mathbf{w}\\cdot f\\mathbf{u}^{0\\perp}\\,\\mathrm{d}x +\n",
    "    g\\beta\\int_\\Omega D^{0}\\nabla\\cdot\\mathbf{w}\\,\\mathrm{d}x, \\\\\n",
    "    R\\lbrack \\phi \\rbrack &= \\int_\\Omega \\phi D^{0}\\,\\mathrm{d}x -\n",
    "    \\beta\\int_\\Omega \\phi\\nabla\\cdot\\mathbf{u}^{0}\\,\\mathrm{d}x.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Here, $\\beta$, $g$, and $H$ are parameters. Within the context of fluid flow on a rotating sphere, $f$ is the Coriolis parameter. $\\mathbf{u}^{0}$ and $D^0$ are previous states of the fields.\n",
    "\n",
    "We require the numerical solution to the saddle-point system:\n",
    "\n",
    "$$\n",
    "    \\mathcal{A}\\mathbf{x} \\equiv\n",
    "    \\begin{bmatrix}\n",
    "        \\color{#800020}{A} & -g\\beta\\color{#2A52BE}{B^T} \\\\\n",
    "        H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C}\n",
    "    \\end{bmatrix}\n",
    "    \\begin{Bmatrix}\n",
    "        U \\\\\n",
    "        D\n",
    "    \\end{Bmatrix} =\n",
    "    \\begin{Bmatrix}\n",
    "        -R\\lbrack \\mathbf{w} \\rbrack \\\\\n",
    "        -R\\lbrack \\phi \\rbrack\n",
    "    \\end{Bmatrix} \\equiv \\mathbf{b}\n",
    "$$\n",
    "\n",
    "for the coefficent vectors $U$ and $D$. A challenge for systems of this type is finding a parameter-indepedent solver for which we have good convergence. First, let's set up an example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example implementation\n",
    "\n",
    "For our domain, let's use a spherical mesh with a radius equal to that of the Earth. We begin by importing Firedrake and defining some relevant constants."
   ]
  },
  {
   "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",
    "\n",
    "R0 = 6371220.0                  # Radius\n",
    "R = Constant(R0)\n",
    "H = Constant(6000.0)            # Mean depth\n",
    "Omega_f = Constant(1.e-4)       # Angular rotation rate\n",
    "g = Constant(10)                # Acceleration due to gravity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create a mesh with quadrilateral elements. We will want to use Firedrake's `CubedSphereMesh`. As with all builtin meshes, we can find out what arguments we need to pass in:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(CubedSphereMesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now we need to specify the number of refinements (say 4). Let's also set the mesh `degree` to be cubic. With this choice of coordinate space, we can better resolve the actual curvature of the sphere using bendy quadrilateral elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = CubedSphereMesh(radius=R0, refinement_level=4, degree=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we just initialize the global normals on this mesh:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = SpatialCoordinate(mesh)\n",
    "mesh.init_cell_orientations(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at the mesh:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake.pyplot import triplot\n",
    "\n",
    "fig = plt.figure()\n",
    "axes = fig.add_subplot(111, projection='3d')\n",
    "triplot(mesh, axes=axes);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define our discrete function spaces. We use a Raviart-Thomas (RT) mixed method on quadrilaterals, taking $V = RTCF_2$ and $U = DQ_1$:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<figure>\n",
    "  <center><img src=\"image/rtcf2_dg1.png\" alt=\"rtcf_dg\" style=\"width: 400px;\"/></center>\n",
    "  <center><figcaption>$RTCF_2$ (left) and $DQ_1 (right)$.\n",
    "      Source: <a href=\"http://femtable.org/\">periodic table of finite elements</a>\n",
    "      </figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = FunctionSpace(mesh, \"RTCF\", 2)\n",
    "U = FunctionSpace(mesh, \"DQ\", 1)\n",
    "W = V * U"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And create an expression for the Coriolis term:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = 2 * Omega_f * x[2] / R"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our initial profiles for $\\mathbf{u}$ and $D$, we set them in steady rotating state:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\mathbf{u}^0 &= \\frac{u_{\\text{max}}}{R}\\left( -y, x, 0\\right), \\\\\n",
    "    D^0 &= H - \\left(R\\Omega_f u_{\\text{max}} + \\frac{u_{\\text{max}}^2}{2}\\right)\\frac{z^2}{gR^2},\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where $u_{\\text{max}} = 40.0 \\text{m}\\text{s}^{-1}$ and $\\Omega_f$ is the planetary rotation rate. In Firedrake, we simply define the expressions in UFL and project/interpolate into the relevant finite element spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_max = Constant(40)\n",
    "u_expr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])\n",
    "D_expr = H -  (R * Omega_f * u_max + u_max**2 / 2.0) * (x[2]**2 / (g * R**2))\n",
    "u0 = Function(V).project(u_expr)\n",
    "D0 = Function(U).interpolate(D_expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need solution functions for $\\mathbf{u}$ and $D$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wh = Function(W, name=\"w_h\")    # Fields in a mixed function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we set the coefficient parameter $\\beta$ and define some test/trial functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = Constant(2000.0)\n",
    "u, D = TrialFunctions(W)\n",
    "w, phi = TestFunctions(W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UFL has all the right objects to specify our problem. However, for the integral:\n",
    "$$\n",
    "\\begin{equation*}\n",
    "\\int_\\Omega \\mathbf{w}\\cdot f\\mathbf{u}^\\perp\\,\\mathrm{d}x,\n",
    "\\end{equation*}\n",
    "$$\n",
    "we need to be a bit careful. UFL does indeed have a `perp` operator, but it does not work for embedded manifolds. That is okay! Since UFL is all in Python, we can create our own `perp` via a Python function which returns a UFL expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "perp = lambda u: cross(CellNormal(mesh), u)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can define our finite element problem in residual form and create a `LinearVariationalProblem`. Here is the problem in residual form:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    0 = F(\\mathbf{u}, D; \\mathbf{w}, \\phi) &=\n",
    "    \\color{#800020}{\\int_\\Omega \\mathbf{w}\\cdot\\mathbf{u}\\,\\mathrm{d}x +\n",
    "    \\beta\\int_\\Omega \\mathbf{w}\\cdot f\\mathbf{u}^\\perp\\,\\mathrm{d}x} \\\\\n",
    "    &- \\color{#2A52BE}{g\\beta\\int_\\Omega D\\nabla\\cdot\\mathbf{w}\\,\\mathrm{d}x} \\\\\n",
    "    &- \\int_\\Omega \\mathbf{w}\\cdot\\mathbf{u}^{0}\\,\\mathrm{d}x\n",
    "    + \\beta\\int_\\Omega \\mathbf{w}\\cdot f\\mathbf{u}^{0\\perp}\\,\\mathrm{d}x \\\\\n",
    "    &- g\\beta\\int_\\Omega D^{0}\\nabla\\cdot\\mathbf{w}\\,\\mathrm{d}x \\\\\n",
    "    &+ \\color{#CC5500}{\\int_\\Omega \\phi D\\,\\mathrm{d}x} +\n",
    "    \\color{#2A52BE}{H\\beta\\int_\\Omega \\phi\\nabla\\cdot\\mathbf{u}\\,\\mathrm{d}x} \\\\\n",
    "    &- \\int_\\Omega \\phi D^{0}\\,\\mathrm{d}x +\n",
    "    \\beta\\int_\\Omega \\phi\\nabla\\cdot\\mathbf{u}^{0}\\,\\mathrm{d}x,\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "and in UFL:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uD_eqn = (inner(w, u) + beta*inner(w, f*perp(u))\n",
    "          - g*beta*div(w)*D\n",
    "          - inner(w, u0) + beta*inner(w, f*perp(u0))\n",
    "          - g*beta*div(w)*D0\n",
    "          + phi*D + H*beta*phi*div(u)\n",
    "          - phi*D0 + H*beta*phi*div(u0))*dx\n",
    "a = lhs(uD_eqn)\n",
    "L = rhs(uD_eqn)\n",
    "\n",
    "uD_problem = LinearVariationalProblem(a, L, wh, constant_jacobian=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## An approximate Schur-complement preconditioner\n",
    "\n",
    "In exact arithmetic, the inverse of the Schur-complement factorization of $\\mathcal{A}$ is:\n",
    "\n",
    "$$\n",
    "   \\mathcal{A}^{-1} =\n",
    "   \\begin{bmatrix}\n",
    "       I & g\\beta \\color{#800020}{A^{-1}}\\color{#2A52BE}{B^T} \\\\\n",
    "       0 & I\n",
    "   \\end{bmatrix}\n",
    "   \\begin{bmatrix}\n",
    "       \\color{#800020}{A^{-1}} & 0 \\\\\n",
    "       0 & \\color{#CC5500}{S}^{-1}\n",
    "   \\end{bmatrix}\n",
    "   \\begin{bmatrix}\n",
    "       I & 0 \\\\\n",
    "       -H\\beta \\color{#2A52BE}{B}\\color{#800020}{A^{-1}} & I\n",
    "   \\end{bmatrix},\n",
    "$$\n",
    "where $\\color{#CC5500}{S}$ is the Schur-complement:\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "    \\color{#CC5500}{S} = \\color{#CC5500}{C} + gH\\beta^2\\color{#2A52BE}{B}\\color{#800020}{A^{-1}}\\color{#2A52BE}{B^T}.\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "Problem: $\\color{#800020}{A^{-1}}$ is _dense_! Instead, let's construct a sparse approximation by using the diagonal of $\\color{#800020}{A}$ to form:\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "    \\color{#CC5500}{\\tilde{S}} = \\color{#CC5500}{C} + gH\\beta^2\\color{#2A52BE}{B}\\text{Diag}(\\color{#800020}{A})^{-1}\\color{#2A52BE}{B^T}.\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we use a block-preconditioner (centered around a full factorization) and approximately invert both $\\color{#800020}{A}$ and $\\color{#CC5500}{\\tilde{S}}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver_parameters = {'ksp_type': 'gmres',\n",
    "                     'ksp_rtol': 1.0e-7,\n",
    "                     'ksp_max_it': 1500,\n",
    "                     'pc_type': 'fieldsplit',\n",
    "                     'pc_fieldsplit': {'type': 'schur',\n",
    "                                       'schur_fact_type': 'full',\n",
    "                                       'schur_precondition': 'selfp'},\n",
    "                     'fieldsplit_0': {'ksp_type': 'preonly',\n",
    "                                      'pc_type': 'bjacobi',\n",
    "                                      'sub_pc_type': 'ilu'},\n",
    "                     'fieldsplit_1': {'ksp_type': 'preonly',\n",
    "                                      'pc_type': 'gamg',\n",
    "                                      'mg_levels': {'ksp_type': 'chebyshev',\n",
    "                                                    'ksp_max_it': 5,\n",
    "                                                    'pc_type': 'bjacobi',\n",
    "                                                    'sub_pc_type': 'ilu'}}}\n",
    "uD_solver_gmres = LinearVariationalSolver(uD_problem,\n",
    "                                          solver_parameters=solver_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we just call solve, store our results, and take a look at the reduction in the residual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()\n",
    "r = assemble(action(a, wh) - L)\n",
    "b = assemble(L)\n",
    "print(\"reduction in residual: %s\" % (r.dat.norm / b.dat.norm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solver convergence\n",
    "\n",
    "Although the code we wrote above works fine, let's take a closer look by inspecting solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gmres_solver_conv(solver):\n",
    "    from firedrake.solving_utils import KSPReasons\n",
    "\n",
    "    print(\"gmres iterations = {}, converged reason = {}\".format(\n",
    "           solver.snes.ksp.getIterationNumber(), \n",
    "           KSPReasons[solver.snes.ksp.getConvergedReason()]))\n",
    "\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay. But what happens if we change the parameters a bit? For example, increasing $\\beta$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(4000.0)\n",
    "wh.assign(0.0)        # re-initialize solution vector\n",
    "uD_solver_gmres.solve()\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(8000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(16000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(32000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(64000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()\n",
    "gmres_solver_conv(uD_solver_gmres)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yikes. Can we do better?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The hybridized mixed method\n",
    "\n",
    "Hybridizing the mixed formulation is an alternative approach which avoids building a global dense operator. We seek approximations $\\hat{\\mathbf{u}}, D, \\lambda \\in \\widehat{V} \\times U \\times T$, where $\\widehat{V}$ is a _discontinuous_ version of $V$ and $T$ is the _trace_ of $V$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$V = RTCF_2$ is continuous in the sense that nodes on the facets shared between cells are topologically identical. The result of rendering $V$ discontinuous removes this association:\n",
    "\n",
    "<ul class=\"list-unstyled list-inline text-center\">\n",
    "  <li>\n",
    "    <img src='image/w2.png' alt='w2' style=\"width: 275px;\"/>\n",
    "    <figcaption>$V$: continuous between facets.</figcaption>\n",
    "  </li>\n",
    "    &#160; &#160; &#160; &#160; &#160; &#160; &#160; &#160;\n",
    "  <li>\n",
    "    <img src='image/w2b.png' alt='w2b' style=\"width: 320px;\"/>\n",
    "    <figcaption>$\\widehat{V}$: no continuity across cell facets.</figcaption>\n",
    "  </li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trace of $V$ is a scalar-valued function space defined only on the mesh skeleton, $\\mathcal{E}$. It is constructed such that function in $T$ belong to the _same_ polynomial space as $\\mathbf{w}\\cdot\\mathbf{n}$, $\\mathbf{w} \\in V$:\n",
    "\n",
    "<figure>\n",
    "    <center><img src='image/w2t.png' alt='w2t' style=\"width: 250px;\"/></center>\n",
    "    <center><figcaption>The space of traces, $T$. Functions in this space are discontinuous across vertices.</figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The approximations $\\hat{\\mathbf{u}}, D, \\lambda$ satisfy the following variational problem:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    \\color{#800020}{\\int_\\Omega \\hat{\\mathbf{w}}\\cdot\\hat{\\mathbf{u}}\\,\\mathrm{d}x +\n",
    "    \\beta\\int_\\Omega \\hat{\\mathbf{w}}\\cdot f\\hat{\\mathbf{u}}^\\perp\\,\\mathrm{d}x}\n",
    "    - \\color{#2A52BE}{g\\beta\\int_\\Omega D\\nabla\\cdot\\hat{\\mathbf{w}}\\,\\mathrm{d}x} +\n",
    "    \\color{#50404D}{\\sum_{K \\in \\Omega} \\int_{\\partial K} \\lambda \\hat{\\mathbf{w}}\\cdot\\mathbf{n}\\,\\mathrm{d}s}\n",
    "    &= -R\\lbrack \\hat{\\mathbf{u}} \\rbrack,\n",
    "    \\quad \\forall \\hat{\\mathbf{w}} \\in \\widehat{V}, \\\\\n",
    "    \\color{#CC5500}{\\int_\\Omega \\phi D\\,\\mathrm{d}x} +\n",
    "    \\color{#2A52BE}{H\\beta\\int_\\Omega \\phi\\nabla\\cdot\\hat{\\mathbf{u}}\\,\\mathrm{d}x} &= -R\\lbrack \\phi \\rbrack,\n",
    "    \\quad \\forall \\phi \\in U, \\\\\n",
    "    \\color{#50404D}{\\sum_{K \\in \\Omega} \\int_{\\partial K} \\gamma \\hat{\\mathbf{u}}\\cdot\\mathbf{n}\\,\\mathrm{d}s}\n",
    "    &= 0, \\quad\\quad\\quad \\forall \\gamma \\in T.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "This is the \"hybridized\" mixed method, resulting in an augmented system with an additional unknown $\\lambda$. This system is solving the same PDE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The discrete hybridized system\n",
    "\n",
    "The resulting discrete system is:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "    \\color{#800020}{\\hat{A}} & -g\\beta\\color{#2A52BE}{B^T} & \\color{#50404D}{E^T} \\\\\n",
    "    H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C} & 0 \\\\\n",
    "    \\color{#50404D}{E} & 0 & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{Bmatrix}\n",
    "    \\hat{U} \\\\\n",
    "    D \\\\\n",
    "    \\Lambda\n",
    "\\end{Bmatrix} =\n",
    "\\begin{Bmatrix}\n",
    "    -R^n\\lbrack \\hat{\\mathbf{w}} \\rbrack \\\\\n",
    "    -R^n\\lbrack \\phi \\rbrack \\\\\n",
    "    0\n",
    "\\end{Bmatrix}.\n",
    "$$\n",
    "\n",
    "Upon initial inspection, it may appear that hybridizing the mixed problem is counter-productive. However, since $\\hat{U}$ and $D$ coupled within the cell interiors only, both can be **eliminated cell-wise** via static condensation!\n",
    "\n",
    "<ul class=\"list-unstyled list-inline text-center\">\n",
    "  <li>\n",
    "    <img src='image/global_mixed_sparsity.png' alt='sparsemat_mixed' style=\"width: 275px;\"/>\n",
    "    <figcaption>Sparsity pattern for the original mixed system.</figcaption>\n",
    "  </li>\n",
    "    &#160; &#160; &#160; &#160; &#160; &#160; &#160; &#160;\n",
    "  <li>\n",
    "    <img src='image/global_hybridized_sparsity.png' alt='sparsemat_mixed' style=\"width: 275px;\"/>\n",
    "    <figcaption>Sparsity pattern for the hybridizable system.</figcaption>\n",
    "  </li>\n",
    "</ul>\n",
    "\n",
    "A reduced problem for $\\Lambda$ can be formed:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "    \\color{#50404D}{E} & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "    \\color{#800020}{\\hat{A}} & -g\\beta\\color{#2A52BE}{B^T}\\\\\n",
    "    H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C}\n",
    "\\end{bmatrix}^{-1}\n",
    "\\begin{bmatrix}\n",
    "    \\color{#50404D}{E^T} \\\\\n",
    "    0\n",
    "\\end{bmatrix}\n",
    "\\Lambda =\n",
    "\\begin{bmatrix}\n",
    "    \\color{#50404D}{E} & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "    \\color{#800020}{\\hat{A}} & -g\\beta\\color{#2A52BE}{B^T}\\\\\n",
    "    H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C}\n",
    "\\end{bmatrix}^{-1}\n",
    "\\begin{Bmatrix}\n",
    "    -R^n\\lbrack \\hat{\\mathbf{w}} \\rbrack \\\\\n",
    "    -R^n\\lbrack \\phi \\rbrack\n",
    "\\end{Bmatrix}.\n",
    "$$\n",
    "\n",
    "There are a number of advantages to inverting the $\\Lambda$ system over the original mixed problem.\n",
    "* The full hybridized system is **never** explicitly assembled;\n",
    "\n",
    "* The block operator:\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "    \\color{#800020}{\\hat{A}} & -g\\beta\\color{#2A52BE}{B^T}\\\\\n",
    "    H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "can be inverted cell-wise.\n",
    "\n",
    "* Once $\\Lambda$ is determined, $\\hat{U}$ and $D$ can be **recovered cell-wise** by inverting the local systems:\n",
    "$$\n",
    "\\begin{Bmatrix}\n",
    "    \\hat{U} \\\\\n",
    "    D\n",
    "\\end{Bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "    \\color{#800020}{\\hat{A}} & -g\\beta\\color{#2A52BE}{B^T}\\\\\n",
    "    H\\beta\\color{#2A52BE}{B} & \\color{#CC5500}{C}\n",
    "\\end{bmatrix}^{-1}\n",
    "\\left(\n",
    "\\begin{Bmatrix}\n",
    "    -R^n\\lbrack \\hat{\\mathbf{w}} \\rbrack \\\\\n",
    "    -R^n\\lbrack \\phi \\rbrack\n",
    "\\end{Bmatrix}\n",
    "-\n",
    "\\begin{bmatrix}\n",
    "    \\color{#50404D}{E^T}\\\\\n",
    "    0\n",
    "\\end{bmatrix}\n",
    "\\Lambda\\right).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "People are often dissuaded from hybridized finite element solvers for their tedious implementation stages. However, Firedrake possesses an abstraction layer for automatically generating element-local dense linear algebra kernels (Slate).\n",
    "In Firedrake, you can write out the dense linear algebra expressions directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Vhat = FunctionSpace(mesh, BrokenElement(V.ufl_element()))\n",
    "Uhat = FunctionSpace(mesh, U.ufl_element())\n",
    "T = FunctionSpace(mesh, FiniteElement(\"HDiv Trace\", mesh.ufl_cell(), 1))\n",
    "Whybrid = Vhat * Uhat * T\n",
    "\n",
    "n = FacetNormal(mesh)\n",
    "uhat, Dhat, lambdar = TrialFunctions(Whybrid)\n",
    "what, phihat, gammar = TestFunctions(Whybrid)\n",
    "\n",
    "# bilinear form (hybrid-mixed system)\n",
    "a_hybrid = ((inner(what, uhat) + beta*inner(what, f*perp(uhat))\n",
    "            - g*beta*div(what)*Dhat\n",
    "            + phihat*Dhat + H*beta*phihat*div(uhat))*dx\n",
    "            + jump(what, n=n)*lambdar('+')*dS\n",
    "            + jump(uhat, n=n)*gammar('+')*dS)\n",
    "\n",
    "# Slate expression for element-wise static condensation\n",
    "AA = Tensor(a_hybrid)\n",
    "A = AA.blocks\n",
    "S = A[2, 0:1] * A[0:1, 0:1].inv * A[0:1, 2]\n",
    "print(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Firedrake's hybridization interface\n",
    "\n",
    "Using Firedrake's symbolic reasoning capability, we can automate the hybridization and static condensation operations using the Python-based preconditioner: `HybridizationPC`. It can be configured through usual PETSc options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hybrid_parameters = {'ksp_type': 'preonly',\n",
    "                     'mat_type': 'matfree',\n",
    "                     'pc_type': 'python',\n",
    "                     'pc_python_type': 'firedrake.HybridizationPC',\n",
    "                     # Solver for the trace system\n",
    "                     'hybridization': {'ksp_type': 'gmres',\n",
    "                                       'pc_type': 'gamg',\n",
    "                                       'pc_gamg_sym_graph': True,\n",
    "                                       'ksp_rtol': 1e-7,\n",
    "                                       'mg_levels': {'ksp_type': 'richardson',\n",
    "                                                     'ksp_max_it': 5,\n",
    "                                                     'pc_type': 'bjacobi',\n",
    "                                                     'sub_pc_type': 'ilu'}}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now one can run the simulation just as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(2000.0)      # set beta back to initial value\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid = LinearVariationalSolver(uD_problem,\n",
    "                                           solver_parameters=hybrid_parameters)\n",
    "uD_solver_hybrid.solve()\n",
    "r = assemble(action(a, wh) - L)\n",
    "b = assemble(L)\n",
    "print(\"reduction in residual: %s\" % (r.dat.norm / b.dat.norm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now how about a closer look at the solver."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solver convergence for the hybridized system\n",
    "\n",
    "Let's stress out the hybridization solver a bit. As before, we write a convenience function for viewing solver convergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hybrid_solver_conv(solver):\n",
    "    from firedrake.solving_utils import KSPReasons\n",
    "\n",
    "    # get the ksp associated with the trace system\n",
    "    trace_ksp = solver.snes.ksp.getPC().getPythonContext().trace_ksp\n",
    "\n",
    "    print(\"gmres iterations (trace sys) = {}, converged reason = {}\".format(\n",
    "          trace_ksp.getIterationNumber(), \n",
    "          KSPReasons[trace_ksp.getConvergedReason()]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as we previously did with our approximate Schur-complement preconditioner, let's go ahead and start a sequence $\\beta$-parametrized runs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(4000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()\n",
    "hybrid_solver_conv(uD_solver_hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(8000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()\n",
    "hybrid_solver_conv(uD_solver_hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(16000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()\n",
    "hybrid_solver_conv(uD_solver_hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(32000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()\n",
    "hybrid_solver_conv(uD_solver_hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(64000.0)\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()\n",
    "hybrid_solver_conv(uD_solver_hybrid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solver time comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set up a problem and solve using preconditioned GMRES and with hybridization. Taking a modest value for $\\beta$, we can use builtin notebook magic to time the execution of the solve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta.assign(32000.0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's time it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "wh.assign(0.0)\n",
    "uD_solver_gmres.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now for hybridization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "wh.assign(0.0)\n",
    "uD_solver_hybrid.solve()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}