05-mixed-poisson.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": [
    "# Systems of more than one variable\n",
    "\n",
    "Almost all the problems that we will want to solve are systems of more than one variable.  For example, the Navier-Stokes equation for fluid flow, $C^0$-conforming discretisations of the Cahn-Hilliard equation for phase separation, or systems that couple different physical processes.\n",
    "\n",
    "Firedrake supports specifying and solving such systems in *mixed* function spaces.  Here we will use a dual formulation of the Poisson equation as a simple introductory example.\n",
    "\n",
    "These systems of equations have a block algebraic form, and we will also look at some of the support for efficient preconditioning strategies that Firedrake offers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem setup\n",
    "Recall the primal form of the Poisson equation.  Find $u$ satisfying\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\nabla^2 u &= -f \\quad\\text{in $\\Omega$}, \\\\\n",
    "   u &= u_0 \\quad\\text{on $\\Gamma_D$},\\\\\n",
    "   \\nabla u \\cdot n &= g \\quad\\text{on $\\Gamma_N$.}\n",
    "   \\end{align}$$\n",
    "   \n",
    "We obtain the dual form by introducing a vector-valued flux $\\sigma = \\nabla u$.  Substituting, the problem becomes.  Find $(\\sigma, u)$ satisfying\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\sigma - \\nabla u &= 0 \\quad \\text{in $\\Omega$}, \\\\\n",
    "\\nabla \\cdot \\sigma &= -f \\quad \\text{in $\\Omega$}, \\\\\n",
    "u &= u_0 \\quad \\text{on $\\Gamma_D$},\\\\\n",
    "\\sigma \\cdot n &= g \\quad\\text{on $\\Gamma \\setminus \\Gamma_D =: \\Gamma_N$.}\n",
    "\\end{align}\n",
    "  $$\n",
    "  \n",
    "For the variational formulation, we introduce a pair of function spaces, $\\Sigma$ and $V$, and seek $(\\sigma, u) \\in \\Sigma \\times V$ such that:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\int_\\Omega (\\sigma \\cdot \\tau + \\nabla \\cdot \\tau u)\\,\\text{d} x &= \n",
    "  \\int_\\Gamma \\tau \\cdot n u\\,\\text{d}s \\quad \\forall \\tau \\in \\Sigma,\\\\ \n",
    "\\int_\\Omega (\\nabla\\cdot\\sigma)v\\,\\text{d} x &= - \\int_\\Omega f v\\,\\text{d}x \n",
    "\\quad \\forall v \\in V.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Notice how the weak condition in the primal form turns into a strong condition on the auxiliary variable $\\sigma$ and the strong condition on $u$ in the primal form appears as a weak condition on $\\sigma$.\n",
    "\n",
    "For this problem we will solve with homogeneous Dirichlet (strong) conditions $u_0 = 0$ on the boundary of a square domain $\\Omega = [0, 1]^2$, and will choose $f = 10\\exp(-100((x-0.5)^2 + (y-0.5)^2))$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As should be familiar by now, we start by importing the Firedrake package and setting up the discrete problem domain."
   ]
  },
  {
   "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 time import time\n",
    "\n",
    "N = 20\n",
    "\n",
    "mesh = UnitSquareMesh(N, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to choose an LBB-stable finite element pair for $W = V\\times\\Sigma$, one choice is the Raviart-Thomas space of degree $k$ for $\\Sigma$ and discontinuous polynomials of degree $k$ for $V$, here we choose $k = 1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rt = FiniteElement(\"Raviart-Thomas\", triangle, 2, variant=\"integral\")\n",
    "Sigma = FunctionSpace(mesh, rt)\n",
    "V = FunctionSpace(mesh, \"DG\", 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that Firedrake inherits the convention from FEniCS that the specified element degree is the degree of the smallest polynomial space that spans the space of the element.  The lowest order Raviart-Thomas element contains all the constants and some linear functions, hence we specify degree 1.\n",
    "\n",
    "Now we build the mixed finite element space $W$, and create test and trial functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = Sigma * V\n",
    "\n",
    "sigma, u = TrialFunctions(W)\n",
    "tau, v = TestFunctions(W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The forcing function $f$ is created as an expression, and we're ready to set up the variational problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = SpatialCoordinate(mesh)\n",
    "f = 10*exp(-100*((x - 0.5)**2 + (y - 0.5)**2))\n",
    "\n",
    "a = dot(sigma, tau)*dx + div(tau)*u*dx + div(sigma)*v*dx\n",
    "L = -f*v*dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need a solution variable, and then we can solve the problem. The problem is symmetric, but indefinite, so we can use the MINRES Krylov method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wh = Function(W)\n",
    "solve(a == L, wh, solver_parameters={\"ksp_type\": \"minres\", \"pc_type\": \"none\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at the solution.\n",
    "Here we're using the matplotlib [subplots](https://matplotlib.org/api/_as_gen/matplotlib.pyplot.subplots.html) function to create a single figure with two subplots oriented in two columns, both of which share the same $x$- and $y$-axes.\n",
    "For each subplot, we're making the aspect ratio of each axis the same, and we're adding a title.\n",
    "Matplotlib can render titles and other labels using [TeX](https://matplotlib.org/tutorials/text/mathtext.html), which we're using to make the $\\sigma$ symbol display nicely in the first title."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "from firedrake.pyplot import quiver, tripcolor\n",
    "\n",
    "sigmah, uh = wh.subfunctions\n",
    "fig, axes = plt.subplots(ncols=2, sharex=True, sharey=True)\n",
    "\n",
    "quiver(sigmah, axes=axes[0])\n",
    "axes[0].set_aspect(\"equal\")\n",
    "axes[0].set_title(\"$\\sigma$\")\n",
    "\n",
    "tripcolor(uh, axes=axes[1])\n",
    "axes[1].set_aspect(\"equal\")\n",
    "axes[1].set_title(\"$u$\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Improving the solver performance\n",
    "\n",
    "Most problems you'll encounter will require effective *preconditioning* to achieve mesh and parameter *independent* convergence.  By this we mean that the number of iterations to obtain the solution is independent of both the mesh refinement, and hopefully also, any coefficient variability.  Here, we have a constant-coefficient problem, so we only need to worry about mesh refinement.\n",
    "\n",
    "If we had a primal Poisson problem, we could treat the issue quite easily using an algebraic multigrid preconditioner such as hypre.  However, for this dual formulation it is slightly more complicated.\n",
    "\n",
    "First, we'll illustrate the issue, then look at some possible solutions.\n",
    "\n",
    "We're going to be running the same problem over and over again at a range of mesh resolutions, so to avoid writing the code again and again we'll wrap the problem up in a function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_solver(N, solver_parameters, aP=None):\n",
    "    mesh = UnitSquareMesh(N, N)\n",
    "    Sigma = FunctionSpace(mesh, rt)\n",
    "    V = FunctionSpace(mesh, \"DG\", 1)\n",
    "    W = Sigma * V\n",
    "\n",
    "    sigma, u = TrialFunctions(W)\n",
    "    tau, v = TestFunctions(W)\n",
    "    x, y = SpatialCoordinate(mesh)\n",
    "    f = 10*exp(-100*((x - 0.5)**2 + (y - 0.5)**2))\n",
    "\n",
    "    a = dot(sigma, tau)*dx + div(tau)*u*dx + div(sigma)*v*dx\n",
    "    L = -f*v*dx\n",
    "    \n",
    "    wh = Function(W)\n",
    "    if aP is not None:\n",
    "        aP = aP(W)\n",
    "    problem = LinearVariationalProblem(a, L, wh, aP=aP)\n",
    "    solver = LinearVariationalSolver(problem, solver_parameters=solver_parameters)\n",
    "    return solver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try solving the problem using MINRES with incomplete Cholesky as a preconditioner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "parameters = {\"ksp_type\": \"minres\", \"mat_type\": \"aij\", \"pc_type\": \"icc\"}\n",
    "from firedrake.solving_utils import KSPReasons\n",
    "for N in [5, 10, 25, 50]:\n",
    "    solver = create_solver(N, parameters)\n",
    "    cpu_timestamp = time()\n",
    "    solver.solve()\n",
    "    cpu_time = time() - cpu_timestamp\n",
    "    print(\"N = {:3d}, iterations = {:3d}, converged reason = {:s}, time = {:.3f}s\".format(\n",
    "        N, solver.snes.ksp.getIterationNumber(), \n",
    "        KSPReasons[solver.snes.ksp.getConvergedReason()],\n",
    "        cpu_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schur complement approaches\n",
    "\n",
    "Many state of the art solution methods for systems of equations rely on block factorisation and block inverses to construct good preconditioners.  Firedrake supports these through PETSc's `fieldsplit` preconditioning type. The mixed Poisson problem\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\int_\\Omega (\\color{#800020}{\\sigma \\cdot \\tau} + \\color{#2A52BE}{\\nabla \\cdot \\tau u})\\,\\text{d} x &= \n",
    "  \\int_\\Gamma \\color{#2A52BE}{\\tau \\cdot n u}\\,\\text{d}s \\quad \\forall \\tau \\in \\Sigma,\\\\ \n",
    "\\int_\\Omega \\color{#BE962A}{(\\nabla\\cdot\\sigma)v}\\,\\text{d} x &= - \\int_\\Omega f v\\,\\text{d}x \n",
    "\\quad \\forall v \\in V.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "is a block system with matrix\n",
    "\n",
    "$$\n",
    "\\mathcal{A} = \\begin{split}\\left(\\begin{matrix} \\color{#800020}A & \\color{#2A52BE}B \\\\\n",
    "\\color{#BE962A}C & 0 \\end{matrix}\\right),\\end{split}\n",
    "$$\n",
    "\n",
    "admitting a factorisation\n",
    "\n",
    "$$\n",
    "\\begin{split}\\left(\\begin{matrix} I & 0 \\\\ \\color{#BE962A}C \\color{#800020}{A^{-1}} & I\\end{matrix}\\right)\n",
    "\\left(\\begin{matrix}\\color{#800020}A & 0 \\\\ 0 & S\\end{matrix}\\right)\n",
    "\\left(\\begin{matrix} I & \\color{#800020}{A^{-1}} \\color{#2A52BE}B \\\\ 0 & I\\end{matrix}\\right),\\end{split}\n",
    "$$\n",
    "\n",
    "with $S = -\\color{#BE962A}C \\color{#800020}{A^{-1}} \\color{#2A52BE}B$ the *Schur complement*.  This has an inverse:\n",
    "\n",
    "$$\n",
    "\\begin{split}\\left(\\begin{matrix} I & -\\color{#800020}{A^{-1}}\\color{#2A52BE}B \\\\ 0 & I \\end{matrix}\\right)\n",
    "\\left(\\begin{matrix} \\color{#800020}{A^{-1}} & 0 \\\\ 0 & S^{-1}\\end{matrix}\\right)\n",
    "\\left(\\begin{matrix} I & 0 \\\\ -\\color{#BE962A}C\\color{#800020}{A^{-1}} & I\\end{matrix}\\right).\\end{split}\n",
    "$$\n",
    "\n",
    "In particular, note that if we drop some of the terms in the factorisation and write\n",
    "\n",
    "$$ \n",
    "\\mathcal{P} = \\begin{split}\\left(\\begin{matrix} \\color{#800020}A & 0 \\\\ 0 & -\\color{#BE962A}C \\color{#800020}{A^{-1}} \\color{#2A52BE}B\\end{matrix}\\right)\\end{split},\n",
    "$$\n",
    "\n",
    "then the system $\\mathcal{P}^{-1}\\mathcal{A}$ has at most four distinct eigenvalues, and so a Krylov method applied to this preconditioned system will converge quickly [Murphy, Golub, and Wathen (2000)](https://doi.org/10.1137/S1064827599355153).\n",
    "\n",
    "This approach therefore reduces the problem of finding an inverse of $\\mathcal{A}$ into the (hopefully simpler) problem of finding good inverses for blocks thereof.\n",
    "\n",
    "Let's try this out on our problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "parameters = {\"ksp_type\": \"minres\", \n",
    "              \"mat_type\": \"aij\", \n",
    "              \"pc_type\": \"fieldsplit\",\n",
    "              \"pc_fieldsplit_type\": \"schur\",\n",
    "              \"pc_fieldsplit_schur_fact_type\": \"diag\",\n",
    "              \"fieldsplit_0_ksp_type\": \"cg\",\n",
    "              \"fieldsplit_0_pc_type\": \"icc\",\n",
    "              \"fieldsplit_0_ksp_rtol\": 1e-12,\n",
    "              \"fieldsplit_1_ksp_type\": \"cg\",\n",
    "              \"fieldsplit_1_pc_type\": \"none\",\n",
    "              \"fieldsplit_1_ksp_rtol\": 1e-12}\n",
    "from firedrake.solving_utils import KSPReasons\n",
    "for N in [5, 10, 25, 50]:\n",
    "    solver = create_solver(N, parameters)\n",
    "    cpu_timestamp = time()\n",
    "    solver.solve()\n",
    "    cpu_time = time() - cpu_timestamp\n",
    "    print(\"N = {:3d}, iterations = {:1d}, converged reason = {:s}, time = {:.3f}s\".format(\n",
    "        N, solver.snes.ksp.getIterationNumber(), \n",
    "        KSPReasons[solver.snes.ksp.getConvergedReason()],\n",
    "        cpu_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great, we got good algorithmic convergence.  Unfortunately, this is still very slow.  This is because the Schur complement operator $S$ is effectively an elliptic operator and we're inverting it with an unpreconditioned Krylov method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PETSc options, an aside\n",
    "\n",
    "Above we saw a larger set of PETSc options than previously.  This warrants a more in depth explanation.  PETSc objects are configurable at *runtime*, and the primary mechanism for controlling their behaviour is the PETSc *options database*.  The \"programming language\" of these options consists of two operations:\n",
    "\n",
    "- Value assignment\n",
    "- string concatenation\n",
    "\n",
    "Each object has an options *prefix* that can be used to distinguish it from other objects (of the same, or different, type).  For example, the linear solve object in PETSc is a `KSP`, and its prefix is `\"ksp_\"`.  Hence, `\"ksp_type\": \"minres\"` says \"Set the type of the KSP object in this solve to minres\".\n",
    "\n",
    "Similarly, the preconditioner is wrapped up in a `PC` object, with prefix `\"pc_\"`.  Here, we say `\"pc_type\": \"fieldsplit\"` which sets the preconditioner type to be `\"fieldsplit\"`.\n",
    "\n",
    "The type of the fieldsplit `PC` is controlled with `\"pc_fieldsplit_type\": \"schur\"`.  This kind of block preconditioner needs inverses for the blocks.  These are provided by KSP objects.  So that we can configure them separately from the main \"outer\" KSP, they have separate prefixes: `\"fieldsplit_0_\"` and `\"fieldsplit_1_\"`.  The former controls the behaviour for $A^{-1}$, the latter for $S^{-1}$.\n",
    "\n",
    "Although this system takes some getting used to, and looks somewhat verbose, it provides a lot of flexibility.  We do not need to modify the code itself, only the options used, to change configuration of the solver."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "Determine how the iterations required to invert the $S$ change when the mesh size is increased.\n",
    "\n",
    "- Hint: The KSPs for the fieldsplit solves are available as `solver.snes.ksp.pc.getFieldSplitSubKSP()` \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Providing an approximation to $S$\n",
    "\n",
    "Our problem is that we haven't provided a preconditioning matrix to invert $S$.  If $A$ is such that $\\text{diag}(A)^{-1}$ is a good approximation to $A^{-1}$, then a sparse approximation $\\tilde{S} = -C \\text{diag}(A)^{-1} B$ is a good approximation to $S$.  PETSc allows us to create this purely using options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "parameters = {\"ksp_type\": \"minres\", \n",
    "              \"mat_type\": \"aij\", \n",
    "              \"pc_type\": \"fieldsplit\",\n",
    "              \"pc_fieldsplit_type\": \"schur\",\n",
    "              \"pc_fieldsplit_schur_fact_type\": \"diag\",\n",
    "              # Construct \\tilde{S} using C diag(A)^{-1} B\n",
    "              \"pc_fieldsplit_schur_precondition\": \"selfp\",\n",
    "              \"fieldsplit_0_ksp_type\": \"cg\",\n",
    "              \"fieldsplit_0_pc_type\": \"icc\",\n",
    "              \"fieldsplit_0_ksp_rtol\": 1e-12,\n",
    "              \"fieldsplit_1_ksp_type\": \"cg\",\n",
    "              \"fieldsplit_1_ksp_converged_reason\": None,\n",
    "              # Use algebraic multigrid on \\tilde{S}\n",
    "              \"fieldsplit_1_pc_type\": \"hypre\",\n",
    "              \"fieldsplit_1_ksp_rtol\": 1e-12}\n",
    "from firedrake.solving_utils import KSPReasons\n",
    "for N in [5, 10, 25, 50, 100]:\n",
    "    solver = create_solver(N, parameters)\n",
    "    cpu_timestamp = time()\n",
    "    solver.solve()\n",
    "    cpu_time = time() - cpu_timestamp\n",
    "    print(\"N = {:3d}, iterations = {:1d}, converged reason = {:s}, time = {:.3f}s\\n\".format(\n",
    "        N, solver.snes.ksp.getIterationNumber(), \n",
    "        KSPReasons[solver.snes.ksp.getConvergedReason()],\n",
    "        cpu_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could now think about reducing the tolerance on the inner solves, reducing the accuracy of the approximation inverses.  For example, approximating $A^{-1}$ and $S^{-1}$ purely by the application of the preconditioner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "parameters = {\"ksp_type\": \"minres\", \n",
    "              \"mat_type\": \"aij\", \n",
    "              \"pc_type\": \"fieldsplit\",\n",
    "              \"pc_fieldsplit_type\": \"schur\",\n",
    "              \"pc_fieldsplit_schur_fact_type\": \"diag\",\n",
    "              # Construct \\tilde{S} using C diag(A)^{-1} B\n",
    "              \"pc_fieldsplit_schur_precondition\": \"selfp\",\n",
    "              \"fieldsplit_0_ksp_type\": \"preonly\",\n",
    "              \"fieldsplit_0_pc_type\": \"icc\",\n",
    "              \"fieldsplit_1_ksp_type\": \"preonly\",\n",
    "              # Use algebraic multigrid on \\tilde{S}\n",
    "              \"fieldsplit_1_pc_type\": \"hypre\"}\n",
    "from firedrake.solving_utils import KSPReasons\n",
    "for N in [5, 10, 25, 50, 100]:\n",
    "    solver = create_solver(N, parameters)\n",
    "    cpu_timestamp = time()\n",
    "    solver.solve()\n",
    "    cpu_time = time() - cpu_timestamp\n",
    "    print(\"N = {:3d}, iterations = {:2d}, converged reason = {:s}, time = {:.3f}s\".format(\n",
    "        N, solver.snes.ksp.getIterationNumber(), \n",
    "        KSPReasons[solver.snes.ksp.getConvergedReason()],\n",
    "        cpu_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Developing efficient solvers for block systems using this approach is quite an experimental science.  The best approach will vary depending on how good the preconditioners you have available for the individual blocks are, how many degrees of freedom the system has, any coefficient variations, the tolerance to which you wish to solve the problem, and probably many more.\n",
    "\n",
    "The strength of Firedrake and PETSc in this set up is the ease with which we can experiment, without needing to recompile code.  For a few more details, and further pointers, you can look at the [extended Firedrake demo on this topic](http://firedrakeproject.org/demos/saddle_point_systems.py.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}