{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2.1.2 Building blocks for programming preconditioners\n",
    "\n",
    "\n",
    "In $\\S$[2.1.1](unit-2.1.1-preconditioners/preconditioner.ipynb), we saw the Jacobi method given as a `local` preconditioner. We now delve deeper into such local preconditioners, which are often useful ingredients/smoothers in more complicated preconditioning strategies. In this tutorial we will see \n",
    "\n",
    "- Jacobi and Gauss-Seidel smoothers/preconditioners,\n",
    "- their block versions, \n",
    "- how to select your own smoothing blocks, and\n",
    "- how to additively combine preconditioners."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from ngsolve.la import EigenValues_Preconditioner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Setup(h=0.1, p=3):\n",
    "    mesh = Mesh(unit_square.GenerateMesh(maxh=h))\n",
    "    fes = H1(mesh, order=p, dirichlet=\"left|bottom\")\n",
    "    u, v = fes.TnT()\n",
    "    a = BilinearForm(grad(u)*grad(v)*dx + u*v*dx)\n",
    "    f = LinearForm(v*dx)\n",
    "    gfu = GridFunction(fes)\n",
    "    return mesh, fes, a, f, gfu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, fes, a, f, gfu = Setup(h=0.1, p=3)\n",
    "a.Assemble()\n",
    "f.Assemble();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Jacobi-preconditioner\n",
    "\n",
    "Let  $A=$ `a.mat` be the assembled matrix, which can be decomposed based on `FreeDofs` ($F$) the remainder ($D$), as in $\\S$[1.3](../unit-1.3-dirichlet/dirichlet.ipynb),\n",
    "\n",
    "$$\n",
    "A = \\left( \\begin{array}{cc} A_{FF} & A_{FD} \\\\ A_{DF} & A_{DD} \\end{array} \\right). \n",
    "$$\n",
    "\n",
    "Then the matrix form of the **point Jacobi preconditioner** is\n",
    "\n",
    "$$\n",
    "J = \\left( \\begin{array}{cc} \\text{diag}(A_{FF})^{-1} & 0  \\\\ 0 & 0  \\end{array} \\right),\n",
    "$$\n",
    "\n",
    "which can be obtained in NGSolve using `CreateSmoother`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preJpoint = a.mat.CreateSmoother(fes.FreeDofs())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NGSolve also gives us a facility to quickly check an estimate of the condition number of the preconditioned matrix by applying the Lanczos algorithm on the preconditioned system.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lams = EigenValues_Preconditioner(mat=a.mat, pre=preJpoint)\n",
    "lams"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An estimate of the condition number $\\kappa = \\lambda_{\\text{max}} / \\lambda_{\\text{min}}$\n",
    "is therefore given as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max(lams)/min(lams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One might wonder if we have gained anything by this point Jacobi preconditioning. What if we did not precondition at all?\n",
    "\n",
    "Not preconditioning is the same as preconditioning by the identity operator on $F$-dofs. One way to realize this identity operator in NGSolve is through the projection into the space of free dofs (i.e., the space spanned by the shape functions corresponding to free dofs). NGSolve provides\n",
    "\n",
    "```\n",
    "Projector(mask, range)   # mask: bit array; range: bool\n",
    "```\n",
    "\n",
    "which projects into the space spanned by the shape functions of the degrees of freedom marked as range in mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preI = Projector(mask=fes.FreeDofs(), range=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note*  that another way to obtain the identity matrix in NGSolve is    \n",
    "```\n",
    "IdentityMatrix(fes.ndof, complex=False).\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lams = EigenValues_Preconditioner(mat=a.mat, pre=preI)\n",
    "max(lams)/min(lams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clearly the point Jacobi preconditioner has reduced the condition number. \n",
    "\n",
    "We can use preconditioners within iterative solvers provided by NGSolve's `solvers` module (which has `MinRes`, `QMR` etc.) Here is an illustration of its use within CG, or the **conjugate gradient** solver: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solvers.CG(mat=a.mat, pre=preJpoint, rhs=f.vec, sol=gfu.vec)\n",
    "Draw(gfu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Gauss-Seidel smoothing\n",
    "\n",
    "The *same* point Jacobi smoother object can also used to perform **point Gauss-Seidel** smoothing. One step of the classical Gauss-Seidel iteration is realized by the method `preJpoint.Smooth()`. It is well known that this iteration converges for matrices like $A$. Below we show how to use it as a linear iterative solver. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu.vec[:] = 0\n",
    "res = f.vec.CreateVector()              # residual \n",
    "projres = f.vec.CreateVector()          # residual projected to freedofs\n",
    "proj = Projector(fes.FreeDofs(), True)\n",
    "\n",
    "for i in range(500):\n",
    "    preJpoint.Smooth(gfu.vec, f.vec)    # one step of point Gauss-Seidel\n",
    "    res.data = f.vec - a.mat*gfu.vec      \n",
    "    projres.data = proj * res\n",
    "    print (\"it#\", i, \", res =\", Norm(projres))\n",
    "Draw (gfu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implement a forward-backward GS preconditioner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *same* point Jacobi smoother object is also able to perform a Gauss-Seidel iteration after reversing the ordering of the points, i.e., a **backward** Gauss-Seidel sweep. One can combine the forward and backward sweeps to construct a symmetric preconditioner, often called the **symmetrized Gauss-Seidel preconditioner**. This offers a good simple illustration of how to construct NGSolve preconditioners from within python. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SymmetricGS(BaseMatrix):\n",
    "    def __init__ (self, smoother):\n",
    "        super(SymmetricGS, self).__init__()\n",
    "        self.smoother = smoother\n",
    "    def Mult (self, x, y):\n",
    "        y[:] = 0.0\n",
    "        self.smoother.Smooth(y, x)\n",
    "        self.smoother.SmoothBack(y,x)\n",
    "    def Height (self):\n",
    "        return self.smoother.height\n",
    "    def Width (self):\n",
    "        return self.smoother.height"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preGS = SymmetricGS(preJpoint)\n",
    "solvers.CG(mat=a.mat, pre=preGS, rhs=f.vec, sol=gfu.vec)\n",
    "Draw(gfu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lams = EigenValues_Preconditioner(mat=a.mat, pre=preGS)\n",
    "max(lams)/min(lams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the condition number now is better than that of the system preconditioned by point Jacobi."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Block Jacobi preconditioner\n",
    "\n",
    "The point Jacobi preconditioner is based on inverses of 1 x 1 diagonal blocks.  Condition numbers can be improved by using larger blocks. It is possible to group dofs into blocks within python and construct an NGSolve preconditioner based on the blocks.\n",
    "\n",
    "Here is an example that constructs vertex-based blocks, using the mesh querying techniques we learnt in $\\S$[1.8](../unit-1.8-meshtopology/meshtopology.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def VertexPatchBlocks(mesh, fes):\n",
    "    blocks = []\n",
    "    freedofs = fes.FreeDofs()\n",
    "    for v in mesh.vertices:\n",
    "        vdofs = set()\n",
    "        for el in mesh[v].elements:\n",
    "            vdofs |= set(d for d in fes.GetDofNrs(el)\n",
    "                         if freedofs[d])\n",
    "        blocks.append(vdofs)\n",
    "    return blocks\n",
    "\n",
    "blocks = VertexPatchBlocks(mesh, fes)\n",
    "print(blocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`CreateBlockSmoother` can now take these blocks and construct a block Jacobi preconditioner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blockjac = a.mat.CreateBlockSmoother(blocks)\n",
    "lams = EigenValues_Preconditioner(mat=a.mat, pre=blockjac)\n",
    "max(lams)/min(lams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multiplicative smoothers and its symmetrized version often yield better condition numbers in practice. We can apply the same code we wrote above for symmetrization (`SymmetricGS`) to the block smoother:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blockgs = SymmetricGS(blockjac)\n",
    "lams = EigenValues_Preconditioner(mat=a.mat, pre=blockgs)\n",
    "max(lams)/min(lams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add a coarse grid correction\n",
    "\n",
    "Dependence of the condition number on degrees of freedom can often be reduced by preconditioners that appropriately use a coarse grid correction. We can experiment with coarse grid corrections using NGSolve's python interface. Here is an example on  how to precondition with a coarse grid correction made using the lowest order subspace of `fes`.\n",
    "\n",
    "In the example below, note that we use `fes.GetDofNrs` again. Previously we used it with argument `el` of type `ElementId`, while now we use it with an argument `v` of type `MeshNode`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def VertexDofs(mesh, fes):\n",
    "    vertexdofs = BitArray(fes.ndof)\n",
    "    vertexdofs[:] = False\n",
    "    for v in mesh.vertices:\n",
    "        for d in fes.GetDofNrs(v):\n",
    "            vertexdofs[d] = True\n",
    "    vertexdofs &= fes.FreeDofs()\n",
    "    return vertexdofs\n",
    "\n",
    "vertexdofs = VertexDofs(mesh, fes)\n",
    "print(vertexdofs)   # bit array, printed 50 chars/line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus we have made a mask `vertexdofs` which reveals all free dofs associated to vertices. If these are labeled $c$ (and the remainder is labeled $f$), then the matrix $A$ can partitioned into \n",
    "\n",
    "$$\n",
    "A = \\left( \\begin{array}{cc} A_{cc} & A_{cf} \\\\ A_{fc} & A_{ff} \\end{array} \\right). \n",
    "$$\n",
    "\n",
    "The matrix `coarsepre` below represents\n",
    "\n",
    "$$\n",
    "\\left( \\begin{array}{cc} A_{cc}^{-1} & 0 \\\\ 0 & 0 \\end{array} \\right). \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coarsepre = a.mat.Inverse(vertexdofs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This matrix can be used for coarse grid correction. \n",
    "\n",
    "##### Pitfall!\n",
    "\n",
    "Note that `coarsepre` is not appropriate as a preconditioner by itself as it has a large null space. You might get the wrong idea from the results of a Lanczos eigenvalue estimation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EigenValues_Preconditioner(mat=a.mat, pre=coarsepre)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But this result only gives the Laczos eigenvalue estimates on the *range* of the preconditioner. The preconditioned operator in this case is simply \n",
    "\n",
    "$$\n",
    "\\left( \\begin{array}{cc} A_{cc}^{-1} & 0 \\\\ 0 & 0 \\end{array} \\right)\n",
    "\\left( \\begin{array}{cc} A_{cc} & A_{cf} \\\\ A_{fc} & A_{ff} \\end{array} \\right)\n",
    " = \n",
    " \\left( \\begin{array}{cc} I  & A_{cc}^{-1} A_{cf} \\\\ 0 & 0  \\end{array} \\right),\n",
    "$$\n",
    "\n",
    "which is a projection into the $c$-dofs. Hence its no surprise that Lanczos estimated the eigenvalues of this operator (on its range) to be just one. But this does not suggest that `coarsepre` has any utility as a preconditioner by itself.\n",
    "\n",
    "##### Additive two-grid preconditioner\n",
    "\n",
    "One well-known  correct way to combine the coarse grid correction with one of the previous smoothers is to combine them additively,  to get an *additive two-grid preconditioner* as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twogrid = coarsepre + blockgs "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This addition of two operators (of type `BaseMatrix`) results in another operator, which is stored as an expression, to be evaluated only when needed.  The 2-grid preconditioner has a very good condition number:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lam = EigenValues_Preconditioner(mat=a.mat, pre=twogrid)\n",
    "lam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combining multigrid with block smoothers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `twogrid` preconditioner becomes more expensive on finer meshes due to the large matrix inversion required for  the computation of `coarsepre`. This can be avoided by replacing `coarsepre` by the multigrid preconditioner we saw in  [2.1.1](../unit-2.1.1-preconditioners/preconditioner.ipynb). It is easy to combine your own block smoothers on the finest grid with the built-in multigrid on coarser levels, as the next example shows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, fes, a, f, gfu = Setup(h=0.5, p=3)\n",
    "Draw(gfu)\n",
    "mg = Preconditioner(a, 'multigrid')  # register mg to a\n",
    "a.Assemble()                         # assemble on coarsest mesh "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us refine a few times to make the problem size larger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(6):   # finer mesh updates & assembly\n",
    "    mesh.Refine() \n",
    "    fes.Update()\n",
    "    a.Assemble()  \n",
    "print('NDOF = ', fes.ndof)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, reset the block smoother to use the blocks on the finest grid. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blocks = VertexPatchBlocks(mesh, fes)\n",
    "blockjac = a.mat.CreateBlockSmoother(blocks)\n",
    "blockgs = SymmetricGS(blockjac)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, add the multigrid and symmetrized block Gauss-Seidel preconditioners together to form a new preconditioner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mgblock = mg.mat + blockgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lam = EigenValues_Preconditioner(mat=a.mat, pre=mgblock)\n",
    "lam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although `mgblock` is similarly conditioned as `twogrid`, it is much more efficient. "
   ]
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}