06_iterative.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "<table>\n",
    " <tr align=left><td><img align=left src=\"https://i.creativecommons.org/l/by/4.0/88x31.png\">\n",
    " <td>Text provided under a Creative Commons Attribution license, CC-BY. All code is made available under the FSF-approved MIT license. (c) Kyle T. Mandli</td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from __future__ import print_function\n",
    "import numpy\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Iterative Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "In this lecture we will consider a number of classical and more modern methods for solving sparse linear systems like those we found from our consideration of boundary value problems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Ways to Solve $A u = f$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We have proposed solving the linear system $A u = f$ which we have implemented naively above with the `numpy.linalg.solve` command but perhaps given the special structure of $A$ here that we can do better.   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Direct Methods (Gaussian Elimination)\n",
    "\n",
    "We could use Gaussian elimination to solve the system (or some factorization) which leads to a solution in a finite number of steps.  For large, sparse methods however these direct solvers are much more expensive in general over iterative solvers.  As was discussed for eigenproblems, iterative solvers start with an initial guess and try to improve on that guess. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Computational Load\n",
    "\n",
    "Now consider using Gaussian elimination on the above matrix.  For good measure let us consider a 3D problem and discretize each dimension with a $N = 100$ leading to $m = 100 \\times 100 \\times 100 = 10^6$ unknowns.\n",
    "\n",
    "Gaussian Elimination - $\\mathcal{O}(m^3)$ operations to solve, $(10^6)^3 = 10^{18}$ operations.\n",
    "\n",
    "Suppose you have a machine that can perform 100 gigaflops (floating point operations per second):\n",
    "$$\n",
    "    \\frac{10^{18}~ [\\text{flop}]}{10^{11}~ [\\text{flop / s}]} = 10^7~\\text{s} \\approx 115~\\text{days}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Memory Load\n",
    "\n",
    "What about memory?\n",
    "\n",
    "We require $m^2$ to store entire array.  In double precision floating point we would require 8-bytes per entry leading to\n",
    "$$\n",
    "    (10^6)^2 ~[\\text{entries}] \\times 8 ~[\\text{bytes / entry}] = 8 \\times 10^{12} ~[\\text{bytes}] = 8 ~[\\text{terabytes}].\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The situation really is not as bad as we are making it out to be as long as we take advantage of the sparse nature of the matrices.  In fact for 1 dimensional problems direct methods can be reduced to $\\mathcal{O}(N)$ in the case for a tridiagonal system.  The situation is not so great for higher-dimensional problems however unless more structure can be leveraged.  Examples of these types of solvers include fast Fourier methods such as fast Poisson solvers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Iterative Methods\n",
    "\n",
    "Iterative methods take a different tact than direct methods.  If we have the system $A x = b$ we form an iterative procedure that applies a function, say $L$, such that\n",
    "$$\n",
    "    \\hat{x}^{(k)} = L(\\hat{x}^{(k-1)})\n",
    "$$\n",
    "where we want errot between the real solution $x$ and $\\hat{x}^{(k)}$ goes to zero as $k \\rightarrow \\infty$.  We will explore these methods in the next lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Jacobi and Gauss-Seidel\n",
    "\n",
    "The Jacobi and Gauss-Seidel methods are simple approaches to introducing an iterative means for solving the problem $Ax = b$ when $A$ is sparse.  Consider again the Poisson problem $u_{xx} = f(x)$ and the finite difference approximation at the point $x_i$\n",
    "$$\n",
    "    \\frac{U_{i-1} - 2 U_i + U_{i+1}}{\\Delta x^2} = f(x_i).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If we rearrange this expression to solve for $U_i$ we have\n",
    "$$\n",
    "    U_i = \\frac{1}{2} (U_{i+1} + U_{i-1}) - f(x_i) \\frac{\\Delta x^2}{2}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For a direct method we would simultaneously find the values of $U_i$, $U_{i+1}$ and $U_{i-1}$ but instead consider the iterative scheme computes an update to the equation above by using the past iterate (values we already know)\n",
    "$$\n",
    "    U_i^{(k+1)} = \\frac{1}{2} (U_{i+1}^{(k)} + U_{i-1}^{(k)}) - f(x_i) \\frac{\\Delta x^2}{2}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Since this allows us to evaluate $U_i^{(k + 1)}$ without knowing the values of $U_{i+1}^{(k)} + U_{i-1}^{(k)}$ we directly evaluate this expression!  This process is called **Jacobi iteration**.  It can be shown that for this particular problem Jacobi iteration will converge from any initial guess $U^{(0)}$ although slowly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Advantages\n",
    " - Matrix $A$ is never stored or created\n",
    " - Storage is optimal\n",
    " - $\\mathcal{O}(m)$ are required per iteration where $m$ is the number of unknowns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Example\n",
    "Let's try to solve the problem before in the BVP section but use Jacobi iterations to replace the direct solve\n",
    "$$\n",
    "    u_{xx} = e^x, \\quad x \\in [0, 1] \\quad \\text{with} \\quad u(0) = 0.0, \\text{ and } u(1) = 3.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "u_a = 0.0\n",
    "u_b = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 100\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Expected iterations needed\n",
    "iterations_J = int(2.0 * numpy.log(delta_x) / numpy.log(1.0 - 0.5 * numpy.pi**2 * delta_x**2))\n",
    "\n",
    "# Solve system\n",
    "# Initial guess for iterations\n",
    "U_new = numpy.zeros(m + 2)\n",
    "U_new[0] = u_a\n",
    "U_new[-1] = u_b\n",
    "convergence_J = numpy.zeros(iterations_J)\n",
    "step_size_J = numpy.zeros(iterations_J)\n",
    "for k in range(iterations_J):\n",
    "    U = U_new.copy()\n",
    "    for i in range(1, m + 1):\n",
    "        U_new[i] = 0.5 * (U[i+1] + U[i-1]) - f(x_bc[i]) * delta_x**2 / 2.0\n",
    "\n",
    "    step_size_J[k] = numpy.linalg.norm(U - U_new, ord=2)\n",
    "    convergence_J[k] = numpy.linalg.norm(u_true(x_bc) - U_new, ord=2)\n",
    "        \n",
    "# Plot result\n",
    "fig = plt.figure()\n",
    "axes = fig.add_subplot(1, 1, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.semilogy(list(range(iterations_J)), step_size_J, 'o')\n",
    "axes.semilogy(list(range(iterations_J)), numpy.ones(iterations_J) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Subsequent Step Size - J\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.semilogy(list(range(iterations_J)), convergence_J, 'o')\n",
    "axes.semilogy(list(range(iterations_J)), numpy.ones(iterations_J) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Convergence to True Solution - J\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "A slight modification to the above leads also to the Gauss-Seidel method.  Programmtically it is easy to see the modification but in the iteration above we now will have\n",
    "$$\n",
    "    U_i^{(k+1)} = \\frac{1}{2} (U_{i+1}^{(k)} + U_{i-1}^{(k+1)}) - f(x_i) \\frac{\\Delta x^2}{2}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "u_a = 0.0\n",
    "u_b = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 100\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Expected iterations needed\n",
    "iterations_GS = int(2.0 * numpy.log(delta_x) / numpy.log(1.0 - numpy.pi**2 * delta_x**2))\n",
    "\n",
    "# Solve system\n",
    "# Initial guess for iterations\n",
    "U = numpy.zeros(m + 2)\n",
    "U[0] = u_a\n",
    "U[-1] = u_b\n",
    "convergence_GS = numpy.zeros(iterations_GS)\n",
    "step_size_GS = numpy.zeros(iterations_GS)\n",
    "success = False\n",
    "for k in range(iterations_GS):\n",
    "    U_old = U.copy()\n",
    "    for i in range(1, m + 1):\n",
    "        U[i] = 0.5 * (U[i+1] + U[i-1]) - f(x_bc[i]) * delta_x**2 / 2.0\n",
    "\n",
    "    convergence_GS[k] = numpy.linalg.norm(u_true(x_bc) - U, ord=2)\n",
    "    step_size_GS[k] = numpy.linalg.norm(U_old - U, ord=2)\n",
    "\n",
    "# Plot result\n",
    "fig = plt.figure()\n",
    "axes = fig.add_subplot(1, 1, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.semilogy(list(range(iterations_GS)), step_size_GS, 'o')\n",
    "axes.semilogy(list(range(iterations_GS)), numpy.ones(iterations_GS) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Subsequent Step Size - GS\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.semilogy(list(range(iterations_GS)), convergence_GS, 'o')\n",
    "axes.semilogy(list(range(iterations_GS)), numpy.ones(iterations_GS) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Convergence to True Solution - GS\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Matrix Splitting Methods\n",
    "\n",
    "One way to view Jacobi and Gauss-Seidel is as a splitting of the matrix $A$ so that\n",
    "$$\n",
    "    A = M - N.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Then the system $A U = b$ can be viewed as\n",
    "$$\n",
    "    M U - N U = b \\Rightarrow MU = NU + b.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Viewing this instead as an iteration we have then\n",
    "$$\n",
    "    M U^{(k+1)} = N U^{(k)} + b.\n",
    "$$\n",
    "The goal then would be to pick $M$ and $N$ such that $M$ contains as much of $A$ as possible while remaining easier to solve than the original system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The resulting update for each of these then becomes\n",
    "$$\n",
    "    U^{(k+1)} = M^{-1} N U^{(k)} + M^{-1} b = G U^{(k)} + c\n",
    "$$\n",
    "where $G$ is called the **iteration matrix** and $c = M^{-1} b$.  We also want\n",
    "$$\n",
    "    u = G u + c\n",
    "$$\n",
    "where $u$ is the true solution of the original $A u = b$, in other words $u$ is the fixed point of the iteration.  Is this fixed point stable though?  If the spectral radius $\\rho(G) < 1$ we can show that in fact the iteration is stable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the similarity between our stability analysis dealing with $||A^{-1}||$ and now $G = M^{-1} N$ which is similar but not identical."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For Jacobi the splitting is \n",
    "$$\n",
    "    M = -\\frac{2}{\\Delta x^2} I, \\quad \\text{and}\\quad N = -\\frac{1}{\\Delta x^2} \\begin{bmatrix}\n",
    "        0 & 1 & \\\\\n",
    "        1 & 0 & 1 \\\\\n",
    "          & \\ddots & \\ddots & \\ddots \\\\\n",
    "          & & 1 & 0 & 1 \\\\\n",
    "          & &   & 1 & 0\n",
    "    \\end{bmatrix}\n",
    "$$\n",
    "(sticking to the Poisson problem).  $M$ is now a diagonal matrix and easy to solve."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For Gauss-Seidel we have\n",
    "$$\n",
    "    M = \\frac{1}{\\Delta x^2} \\begin{bmatrix}\n",
    "        -2 &  & \\\\\n",
    "         1 & -2 &  \\\\\n",
    "           & \\ddots & \\ddots \\\\\n",
    "           & & 1 & -2 & \\\\\n",
    "           & &   & 1 & -2\n",
    "    \\end{bmatrix} \\quad \\text{and} \\quad \n",
    "    N = -\\frac{1}{\\Delta x^2} \\begin{bmatrix}\n",
    "         0 & 1 & \\\\\n",
    "          & 0 & 1 \\\\\n",
    "          & & \\ddots & \\ddots \\\\\n",
    "           & &  & 0 & 1\\\\\n",
    "           & &   &  & 0\n",
    "    \\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Stopping Criteria\n",
    "\n",
    "How many iterations should we perform?  Let $E^{(k)}$ represent the error present at step $k$.  If we want to reduce the error at the first step $E^{(0)}$ to order $\\epsilon$ then we have\n",
    "$$\n",
    "    ||E^{(k)}|| \\approx \\epsilon ||E^{(0)}||.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Under suitable assumption we can bound the error in the 2-norm as\n",
    "$$\n",
    "    ||E^{(k)}||_2 \\leq \\rho(G)^k ||E^{(0)}||_2.\n",
    "$$\n",
    "where $\\rho(G)$ is the spectral radius of the iteration matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Moving back to our estimate of the number of iterations we can combine our two expressions involving the error $E$ by taking $\\Delta x \\rightarrow 0$ which allows us to write\n",
    "$$\n",
    "    k \\approx \\frac{\\log \\epsilon}{\\log \\rho(G)}\n",
    "$$\n",
    "taking into account error convergence.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Picking $\\epsilon$ is a bit tricky but one natural criteria to use would be $\\epsilon = \\mathcal{O}(\\Delta x^2)$ since our original discretization was 2nd-order accurate.  This leads to\n",
    "$$\n",
    "    k = \\frac{2 \\log \\Delta x}{\\log \\rho}.\n",
    "$$\n",
    "This also allows us to estimate the total number of operations that need to be used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For Jacobi we have the spectral radius of $G$ as\n",
    "$$\n",
    "    \\rho_J \\approx 1 - \\frac{1}{2} \\pi^2 \\Delta x^2.\n",
    "$$\n",
    "so that\n",
    "$$\n",
    "    k = \\mathcal{O}(m^2 \\log m) \\quad \\text{as} \\quad m \\rightarrow \\infty.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Combining this with the previous operation count per iteration we find that Jacobi would lead to $\\mathcal{O}(m^3 \\log m)$ work which is not very promising.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For two dimensions we have $\\mathcal{O}(m^4 \\log m)$ so even compared to Gaussian elimination this approach is not ideal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What about Gauss-Seidel?  Here the spectral radius is approximately\n",
    "$$\n",
    "    \\rho_{GS} \\approx 1 - \\pi^2 \\Delta x^2\n",
    "$$\n",
    "so that\n",
    "$$\n",
    "    k = \\frac{2 \\times \\log \\Delta x}{\\log (1 - \\pi^2 \\Delta x^2)}\n",
    "$$\n",
    "which still does not lead to any advantage over direct solvers.  It does show that Gauss-Seidel does actually converge faster due to the factor of 2 difference between $\\rho_J$ and $\\rho_{GS}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Successive Overrelaxation (SOR)\n",
    "\n",
    "Well that's a bit dissapointing isn't it?  These iterative schemes do not seem to be worth much but it turns out we can do better with a slight modification to Gauss-Seidel.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If you look at Gauss-Seidel iteration it turns out it moves $U$ in the correct direction to $u$ but is very conservative in the amount.  If instead we do\n",
    "$$\\begin{aligned}\n",
    "    U^{GS}_i &= \\frac{1}{2} \\left(U^{(k+1)}_{i-1} + U^{(k)}_{i+1} - \\Delta x^2 f_i\\right) \\\\\n",
    "    U^{(k+1)}_i &= U_i^{(k)} + \\omega \\left( U_i^{GS} - U_i^{(k)}\\right )\n",
    "\\end{aligned}$$\n",
    "where we get to pick $\\omega$ we can do much better.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If $\\omega = 1$ then we get back Gauss-Seidel.  \n",
    "\n",
    "If $\\omega < 1$ we move even less and converges even more slowly (although is sometimes used for multigrid under the name **underrelaxation**).  \n",
    "\n",
    "If $\\omega > 1$ then we move further than Gauss-Seidel suggests and any method where $\\omega > 1$ is known as **successive overrelaxation** (SOR)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can write this as a matrix splitting method as well.  We can combine the two-step formula above to find\n",
    "$$\n",
    "    U^{(k+1)}_i = \\frac{\\omega}{2} \\left( U^{(k+1)}_{i-1} + U^{(k)}_{i+1} - \\Delta x^2 f_i \\right ) + (1 - \\omega) U_i^{(k)}\n",
    "$$\n",
    "corresponding to a matrix splitting of\n",
    "$$\n",
    "    M = \\frac{1}{\\omega} (D - \\omega L) \\quad \\text{and} \\quad N = \\frac{1}{\\omega} ((1-\\omega) D + \\omega U)\n",
    "$$\n",
    "where $D$ is the diagonal of the matrix $A$, and $L$ and $U$ are the lower and upper triangular parts without the diagonal of $A$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It can be shown that picking an $\\omega$ such that $0 < \\omega < 2$ the SOR method converges.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It turns out we can also find an optimal $\\omega$ for a wide class of problems.  For Poisson problems in any number of space dimensions for instance it can be shown that SOR method converges optimaly if\n",
    "$$\n",
    "    \\omega_{opt} = \\frac{2}{1 + \\sin(\\pi \\Delta x)} \\approx 2 - 2 \\pi \\Delta x.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What about the number of iterations?  We can follow the same tactic as before with the spectral radius of $G_{SOR}$ now\n",
    "$$\n",
    "    \\rho = \\omega_{opt} - 1 \\approx 1 - 2 \\pi \\Delta x.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This leads to an iteration count of\n",
    "$$\n",
    "    k = \\mathcal{O}(m \\log m)\n",
    "$$\n",
    "an order of magnitude better than Gauss-Seidel alone!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "u_a = 0.0\n",
    "u_b = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 50\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# SOR parameter\n",
    "omega = 2.0 / (1.0 + numpy.sin(numpy.pi * delta_x))\n",
    "\n",
    "# Expected iterations needed\n",
    "iterations_SOR = int(2.0 * numpy.log(delta_x) / numpy.log(1.0 - 2.0 * numpy.pi * delta_x)) * 2\n",
    "\n",
    "# Solve system\n",
    "# Initial guess for iterations\n",
    "U = numpy.zeros(m + 2)\n",
    "U[0] = u_a\n",
    "U[-1] = u_b\n",
    "step_size_SOR = numpy.zeros(iterations_SOR)\n",
    "convergence_SOR = numpy.zeros(iterations_SOR)\n",
    "for k in range(iterations_SOR):\n",
    "    U_old = U.copy()\n",
    "    for i in range(1, m + 1):\n",
    "        U_gs = 0.5 * (U[i-1] + U[i+1] - delta_x**2 * f(x_bc[i]))\n",
    "        U[i] += omega * (U_gs - U[i])\n",
    "\n",
    "    step_size_SOR[k] = numpy.linalg.norm(U_old - U, ord=2)\n",
    "    convergence_SOR[k] = numpy.linalg.norm(u_true(x_bc) - U, ord=2)\n",
    "        \n",
    "# Plot result\n",
    "fig = plt.figure()\n",
    "axes = fig.add_subplot(1, 1, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.semilogy(list(range(iterations_SOR)), step_size_SOR, 'o')\n",
    "axes.semilogy(list(range(iterations_SOR)), numpy.ones(iterations_SOR) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Subsequent Step Size - SOR\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.semilogy(list(range(iterations_SOR)), convergence_SOR, 'o')\n",
    "axes.semilogy(list(range(iterations_SOR)), numpy.ones(iterations_SOR) * delta_x**2, 'r--')\n",
    "axes.set_title(\"Convergence to True Solution - SOR\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Plotting all the convergence rates\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.semilogy(range(iterations_J), step_size_J, 'r', label=\"Jacobi\")\n",
    "axes.semilogy(range(iterations_GS),  step_size_GS, 'b', label=\"Gauss-Seidel\")\n",
    "axes.semilogy(range(iterations_SOR),  step_size_SOR, 'k', label=\"SOR\")\n",
    "axes.semilogy(range(iterations_J), numpy.ones(iterations_J) * delta_x**2, 'r--')\n",
    "axes.legend(loc=1)\n",
    "axes.set_title(\"Comparison of Step Sizes\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.semilogy(range(iterations_J), convergence_J, 'r', label=\"Jacobi\")\n",
    "axes.semilogy(range(iterations_GS),  convergence_GS, 'b', label=\"Gauss-Seidel\")\n",
    "axes.semilogy(range(iterations_SOR),  convergence_SOR, 'k', label=\"SOR\")\n",
    "axes.semilogy(range(iterations_J), numpy.ones(iterations_J) * delta_x**2, 'r--')\n",
    "axes.legend(loc=1)\n",
    "axes.set_title(\"Comparison of Convergence Rates\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.semilogy(range(iterations_SOR), step_size_J[:iterations_SOR], 'r', label=\"Jacobi\")\n",
    "axes.semilogy(range(iterations_SOR),  step_size_GS[:iterations_SOR], 'b', label=\"Gauss-Seidel\")\n",
    "axes.semilogy(range(iterations_SOR),  step_size_SOR, 'k', label=\"SOR\")\n",
    "axes.semilogy(range(iterations_SOR), numpy.ones(iterations_SOR) * delta_x**2, 'r--')\n",
    "axes.legend(loc=1)\n",
    "axes.set_title(\"Comparison of Step Sizes\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.semilogy(range(iterations_SOR), convergence_J[:iterations_SOR], 'r', label=\"Jacobi\")\n",
    "axes.semilogy(range(iterations_SOR),  convergence_GS[:iterations_SOR], 'b', label=\"Gauss-Seidel\")\n",
    "axes.semilogy(range(iterations_SOR),  convergence_SOR, 'k', label=\"SOR\")\n",
    "axes.semilogy(range(iterations_SOR), numpy.ones(iterations_SOR) * delta_x**2, 'r--')\n",
    "axes.legend(loc=3)\n",
    "axes.set_title(\"Comparison of Convergence Rates\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Descent Methods\n",
    "\n",
    "One special case of matrices are amenable to another powerful way to iterate to the solution.  A matrix is said to be **symmetric positive definite** (SPD) if \n",
    "$$\n",
    "    x^T A x > 0 \\quad \\forall \\quad x \\neq 0.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Check to see if \n",
    "$$\n",
    "    A = \\begin{bmatrix}\n",
    "        2 &-1 &0 &0 \\\\\n",
    "        -1 & 2 & -1 & 0 \\\\\n",
    "        0 & -1 & 2 & -1 \\\\\n",
    "        0 & 0 & -1 & 2\n",
    "    \\end{bmatrix}\n",
    "$$\n",
    "is symmetric positive definite."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now define a function $\\phi: \\mathbb R^m \\rightarrow \\mathbb R$ such that\n",
    "$$\n",
    "    \\phi(u) = \\frac{1}{2} u^T A u - u^T f.\n",
    "$$\n",
    "This is a quadratic function in the variables $u_i$ and in the case where $m = 2$ forms a parabolic bowl.  Since this is a quadratic function there is a unique minimum, $u^\\ast$.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Lets see how approaching the problem like this helps us:\n",
    "\n",
    "For the $m = 2$ case write the function $\\phi(u)$ out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "$$\n",
    "    \\phi(u) = \\frac{1}{2} (A_{11} u_1^2 + A_{12} u_1 u_2 + A_{21} u_1 u_2 + A_{22} u^2_2) - u_1 f_1 - u_2 f_2\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What property of the matrix $A$ simplifies the expression above?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "   Symmetry!  This implies that $A_{21} = A_{12}$ and the expression above simplifies to \n",
    "   $$\n",
    "       \\phi(u) = \\frac{1}{2} (A_{11} u_1^2 + 2 A_{12} u_1 u_2 + A_{22} u^2_2) - u_1 f_1 - u_2 f_2\n",
    "   $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now write two expressions that when evaluated at $u^\\ast$ are identically 0 that express that $u^\\ast$ minimizes $\\phi(u)$.\n",
    "\n",
    "   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Since $u^\\ast$ minimizes $\\phi(u)$ we know that the first derivatives should be zero at the minimum:\n",
    "   $$\\begin{aligned}\n",
    "       \\frac{\\partial \\phi}{\\partial u_1} &= A_{11} u_1 + A_{12} u_2 - f_1 = 0 \\\\\n",
    "       \\frac{\\partial \\phi}{\\partial u_1} &= A_{21} u_1 + A_{22} u_2 - f_2 = 0\n",
    "   \\end{aligned}$$\n",
    "   Note that these equations can be rewritten as\n",
    "   $$\n",
    "       A u = f.\n",
    "   $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Therefore $\\min \\phi$ is equivalent to solving $A u = f$!\n",
    "\n",
    "This is a common type of reformulation for many problems where it may be easier to treat a given equation as a minimization problem rather than directly solve it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Note that this is not quite the matrix that we have been using for our Poisson problem so far which is actually symmetric negative definite although these same methods work as well.  In this case we actually want to find the maximum of $\\phi$ instead, other than that everything is the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Also note that if $A$ is indefinite then the eigenvalues of $A$ will change sign and instead of a stable minimum or maximum we have a saddle point which are much more difficult to handle (GMRES can for instance)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Method of Steepest Descent\n",
    "\n",
    "So now we turn to finding the $u^\\ast$ that minimizes the function $\\phi(u)$.  The simplest approach to this is called the **method of steepest descent** which finds the direction of the largest gradient of $\\phi(u)$ and goes in that direction.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Mathematically we then have\n",
    "$$\n",
    "    u^{(k+1)} = u^{(k)} - \\alpha^{(k)} \\nabla \\phi(u^{(k)})\n",
    "$$\n",
    "where $\\alpha^{(k)}$ will be the step size chosen in the direction we want to go.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can find $\\alpha$ by\n",
    "$$\n",
    "    \\alpha^{(k)} = \\min_{\\alpha \\in \\mathbb R} \\phi\\left(u^{(k)} - \\alpha \\nabla \\phi(u^{(k)}\\right),\n",
    "$$\n",
    "i.e. the $\\alpha$ that takes us just far enough so that if we went any further $\\phi$ would increase.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This implies that $\\alpha^{(k)} \\geq 0$ and $\\alpha^{(k)} = 0$ only if we are at the minimum of $\\phi$.  We can compute the gradient of $\\phi$ as\n",
    "$$\n",
    "    \\nabla \\phi(u^{(k)}) = A u^{(k)} - f \\equiv -r^{(k)}\n",
    "$$\n",
    "where $r^{(k)}$ is the residual defined as\n",
    "$$\n",
    "    r^{(k)} = f - A u^{(k)}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Looking back at the definition of $\\alpha^{(k)}$ then leads to the conclusion that the $\\alpha$ that would minimize the expression would be the one that satisfies\n",
    "$$\n",
    "    \\frac{\\text{d} \\phi(\\alpha)}{\\text{d} \\alpha} = 0. \n",
    "$$\n",
    "\n",
    "To find this note that\n",
    "$$\n",
    "    \\phi(u + \\alpha r) = \\left(\\frac{1}{2} u^T A u - u^T f \\right) + \\alpha(r^T A u - r^T f) + \\frac{1}{2} \\alpha^2 r^T A r\n",
    "$$\n",
    "so that the derivative becomes\n",
    "$$\n",
    "    \\frac{\\text{d} \\phi(\\alpha)}{\\text{d} \\alpha} = r^T A u - r^T f + \\alpha r^T A r\n",
    "$$\n",
    "\n",
    "Setting this to zero than leads to\n",
    "$$\n",
    "    \\alpha = \\frac{r^T r}{r^T A r}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "u_a = 0.0\n",
    "u_b = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 50\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Construct matrix A\n",
    "A = numpy.zeros((m, m))\n",
    "diagonal = numpy.ones(m) / delta_x**2\n",
    "A += numpy.diag(diagonal * 2.0, 0)\n",
    "A += numpy.diag(-diagonal[:-1], 1)\n",
    "A += numpy.diag(-diagonal[:-1], -1)\n",
    "\n",
    "# Construct right hand side\n",
    "b = -f(x)\n",
    "b[0] += u_a / delta_x**2\n",
    "b[-1] += u_b / delta_x**2\n",
    "\n",
    "# Algorithm parameters\n",
    "MAX_ITERATIONS = 10000\n",
    "tolerance = 1e-3\n",
    "\n",
    "# Solve system\n",
    "U = numpy.empty(m)\n",
    "convergence_SD = numpy.zeros(MAX_ITERATIONS)\n",
    "step_size_SD = numpy.zeros(MAX_ITERATIONS)\n",
    "success = False\n",
    "for k in range(MAX_ITERATIONS):\n",
    "    r = b - numpy.dot(A, U)\n",
    "    if numpy.linalg.norm(r, ord=2) < tolerance:\n",
    "        success = True\n",
    "        break\n",
    "        \n",
    "    alpha = numpy.dot(r, r) / numpy.dot(r, numpy.dot(A, r))\n",
    "    U = U + alpha * r\n",
    "\n",
    "    step_size_SD[k] = numpy.linalg.norm(alpha * r, ord=2)\n",
    "    convergence_SD[k] = numpy.linalg.norm(u_true(x) - U, ord=2)\n",
    "        \n",
    "if not success:\n",
    "    print(\"Iteration failed to converge!\")\n",
    "    print(convergence_SD[-1])\n",
    "else:\n",
    "    # Plot result\n",
    "    fig = plt.figure()\n",
    "    axes = fig.add_subplot(1, 1, 1)\n",
    "    axes.plot(x, U, 'o', label=\"Computed\")\n",
    "    axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "    axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "    axes.set_xlabel(\"x\")\n",
    "    axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "    fig = plt.figure()\n",
    "    fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "    \n",
    "    axes = fig.add_subplot(1, 2, 1)\n",
    "    axes.semilogy(list(range(k)), step_size_SD[:k], 'o')\n",
    "    axes.semilogy(list(range(k)), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Step Size\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "    axes = fig.add_subplot(1, 2, 2)\n",
    "    axes.semilogy(list(range(k)), convergence_SD[:k], 'o')\n",
    "    axes.semilogy(list(range(k)), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Convergence\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||u(x) - U^{(k-1)}||_2$\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "#### Convergence of Steepest Descent\n",
    "\n",
    "What controls the convergence of steepest descent?  It turns out that the shape of the parabolic bowl formed by $\\phi$ is the major factor determining the convergence of steepest descent.  \n",
    "\n",
    "For example, if $A$ is a scalar multiple of the identity than these ellipses are actually circles and steepest descent converges in $m$ steps.  If $A$ does not lead to circles, the convergence is based on the ratio between the semi-major and semi-minor axis of the resulting ellipses $m$ dimensional ellipses.  \n",
    "\n",
    "This is controlled by the smallest and largest eigenvalues of the matrix $A$ hence why steepest descent grows increasingly difficult as $m$ increases for the Poisson problem.  Note that this also relates to the condition number of the matrix in the $\\ell_2$ norm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "![Ellipses](./images/ellipses.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def steepest_descent(A, U, b, axes):\n",
    "    MAX_ITERATIONS = 10000\n",
    "    tolerance = 1e-3\n",
    "    success = False\n",
    "    iteration_locations = []\n",
    "    for k in range(MAX_ITERATIONS):\n",
    "        axes.text(U[0] + 0.1, U[1] + 0.1, str(k), fontsize=12)\n",
    "        axes.plot(U[0], U[1], 'ro')\n",
    "        iteration_locations.append(U)\n",
    "        r = b - numpy.dot(A, U)\n",
    "        if numpy.linalg.norm(r, ord=2) < tolerance:\n",
    "            success = True\n",
    "            break\n",
    "\n",
    "        alpha = numpy.dot(r, r) / numpy.dot(r, numpy.dot(A, r))\n",
    "        U = U + alpha * r\n",
    "        \n",
    "    if success:\n",
    "        return k, iteration_locations\n",
    "    else:\n",
    "        raise Exception(\"Iteration did not converge.\")\n",
    "\n",
    "phi = lambda X, Y, A: 0.5 * (A[0, 0] * X**2 + A[0, 1] * X * Y + A[1, 0] * X * Y + A[1, 1] * Y**2) - X * f[0] - Y * f[1]\n",
    "\n",
    "x = numpy.linspace(-15, 15, 100)\n",
    "y = numpy.linspace(-15, 15, 100)\n",
    "X, Y = numpy.meshgrid(x, y)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.identity(2)\n",
    "f = numpy.array([0.0, 0.0])\n",
    "k, iteration_locations = steepest_descent(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "phi_levels = numpy.array([phi(U[0], U[1], A) for U in iteration_locations])\n",
    "phi_levels.sort()\n",
    "axes.contour(X, Y, phi(X, Y, A), levels=phi_levels, colors='r')\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.array([[2, -1], [-1, 2]])\n",
    "f = numpy.array([1.0, 2.0])\n",
    "k, iteration_locations = steepest_descent(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "phi_levels = numpy.array([phi(U[0], U[1], A) for U in iteration_locations])\n",
    "phi_levels.sort()\n",
    "axes.contour(X, Y, phi(X, Y, A), levels=phi_levels, colors='r')\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.array([[2, -1.8], [-1.8, 2]])\n",
    "f = numpy.array([1.0, 2.0])\n",
    "k, iteration_locations = steepest_descent(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "phi_levels = numpy.array([phi(U[0], U[1], A) for U in iteration_locations])\n",
    "phi_levels.sort()\n",
    "axes.contour(X, Y, phi(X, Y, A), levels=phi_levels, colors='r')\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Each of these ellipses is related to the eigenstructure of $A$ such that\n",
    "$$\n",
    "    A v_j - f = \\lambda_j (v_j - u^\\ast)\n",
    "$$\n",
    "for some $\\lambda_j$.  Knowing that $A u^\\ast = f$ leads to\n",
    "$$\n",
    "    A (v_j - u^\\ast) = \\lambda_j (v_j - u^\\ast)\n",
    "$$\n",
    "therefore $v_j - u^\\ast$ form the eigenvectors of the matrix $A$ with corresponding eigenvalues $\\lambda_j$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If a particular set of $\\lambda_j$s are not distinct than the ellipse is in fact a circle, demonstrating that any direction pointing to $u^\\ast$ in this direction is an eigenvector (non-unique in this sub-space)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can also relate this eigenstructure and geometric arguments to the matrix's condition number $\\kappa$.  Let $v_1$ and $v_2$ be vectors that lie along the curve $\\phi(u) = 1$, we then have\n",
    "$$\n",
    "    \\frac{1}{2} v^T_j A v_j - v_j^T A u^\\ast = 1.\n",
    "$$\n",
    "Combining this expression with our previous eigenvector expression and taking the inner-product with the eigenvector $v_j - u^\\ast$ leads to\n",
    "$$\n",
    "    ||v_j - u^\\ast||^2_2 = \\frac{2 + (u^\\ast)^T A u^\\ast}{\\lambda_j}.\n",
    "$$\n",
    "Now turning back to $v_1$ and $v_2$ we have their ratios as\n",
    "$$\n",
    "    \\frac{||v_1 - u^\\ast||_2}{||v_2 - u^\\ast||_2} = \\sqrt{\\frac{\\lambda_2}{\\lambda_1}} = \\sqrt{\\kappa_2(A)}.\n",
    "$$\n",
    "This last expression tells us that the more ellipsoidal these sub-spaces are the more difficult it will be to solve $A u^\\ast = f$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Projection Interpretation\n",
    "\n",
    "One way to interpret the method of steepest descent is as an iterative projection method.  \n",
    "\n",
    "Say we want to solve $A x = b$, $A \\in \\mathbb R^{m \\times m}$ and $b \\in \\mathbb R^m$.  \n",
    "\n",
    "Let us suppose we have two subspaces in $\\mathbb R^m$, $\\mathcal K$ the search subspace, and $\\mathcal L$ the subspace of constraints.  The condition\n",
    "$$\n",
    "    b - A x \\perp \\mathcal{L}\n",
    "$$\n",
    "suggests that the residual vector $b - A x$ is orthogonal to the $\\mathcal{L}$ subspace."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "But how does $\\mathcal K$ fit into this?  Modify the original statement so that\n",
    "$$\n",
    "    \\tilde{x~} \\in \\mathcal{K}\n",
    "$$\n",
    "and then\n",
    "$$\n",
    "    b - A \\tilde{x~} \\perp \\mathcal{L}.\n",
    "$$\n",
    "These are known as *Petrov-Galerkin* conditions.\n",
    "\n",
    "If $\\mathcal{K} = \\mathcal{L}$ then this is an orthogonal projection."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Turning this into an iterative method let our initial guess be $x^{(0)}$, then\n",
    "$$\n",
    "    \\tilde{x~} \\in x^{(0)} + \\mathcal{K}\n",
    "$$\n",
    "such that\n",
    "$$\n",
    "    b - A \\tilde{x~} \\perp \\mathcal{L}.\n",
    "$$\n",
    "Note that this implies that we are expanding the search space by the part of $\\text{span}(x^{(0)})$ that is not included in $\\mathcal{K}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can rewrite this in a more suggestive form by letting\n",
    "$$\n",
    "    \\tilde{x~} = x^{(0)} + \\delta\n",
    "$$\n",
    "where $\\delta \\in \\mathcal{K}$ is some vector that will be our step.  Also assign the residual vector an index so that we have\n",
    "$$\n",
    "    r^{(k)} = b - A x^{(k)}\n",
    "$$\n",
    "at the $k$th step.\n",
    "\n",
    "We then formulate a new problem such that\n",
    "$$\n",
    "    r^{(0)} - A \\delta \\perp \\mathcal{L}\n",
    "$$\n",
    "leading to an iteration like statement:\n",
    "$$\\begin{aligned}\n",
    "    \\tilde{x~} &= x^{(0)} + \\delta, & \\quad \\quad & \\delta \\in \\mathcal{K} \\\\\n",
    "    (r^{(0)} - A \\delta) \\cdot w &= 0 & \\quad \\quad & \\forall w \\in \\mathcal{L}\n",
    "\\end{aligned}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Example:  One-Dimensional Projection\n",
    "\n",
    "One-dimensional projection methods simply construct $\\mathcal{K}$ and $\\mathcal{L}$ such that they are\n",
    "$$\n",
    "    \\mathcal{K} = \\text{span}(v) \\quad\\quad \\mathcal{L} = \\text{span}(w)\n",
    "$$\n",
    "for two vectors $v$ and $w$.  In this case we can write the projection update as\n",
    "$$\n",
    "    x^{(k+1)} = x^{(k)} + \\alpha r\n",
    "$$\n",
    "where $r$ is the residual and \n",
    "$$\n",
    "    \\alpha = \\frac{r \\cdot w}{A v \\cdot w}.\n",
    "$$\n",
    "\n",
    "In the case of steepest descent $v = w = r$ noting that these lead to orthogonal projections as $\\mathcal{K} = \\mathcal{L}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Generalizing this a bit we can think of the steepest descent method in terms of projections.  Let $\\mathcal{K} = \\mathcal{L}$ and again assume that $A$ is symmetric positive-definite.  Define the error at the $k$th step as\n",
    "$$\n",
    "    E^{(k)} = x^\\ast - x^{(k)}.\n",
    "$$\n",
    "\n",
    "In this case consider a single step from $k=0$ to $k=1$.  Then we would have\n",
    "$$\n",
    "    r^{(1)} = b - A (x^{(0)} + \\delta) = r^{(0)} - A \\delta\n",
    "$$\n",
    "and\n",
    "$$\n",
    "    A E^{(1)} = r^{(1)} = A(E^{(0)} - \\delta)\n",
    "$$\n",
    "where $\\delta$ is a result of the projection so that\n",
    "$$ \n",
    "    (r^{(0)} - A \\delta) \\cdot w = 0 \\quad \\quad \\forall w \\in \\mathcal{K}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### $A$-Conjugate Search Directions and Conjugate Gradient\n",
    "\n",
    "An alternative to steepest descent is to choose a slightly different direction to descend down.  Generalizing our step from above let the iterative scheme be  \n",
    "$$\n",
    "    u^{(k+1)} = u^{(k)} - \\alpha^{(k)} p^{(k)}\n",
    "$$\n",
    "and as before we want to pick an $\\alpha$ such that\n",
    "$$\n",
    "    \\min_{\\alpha} \\phi(u^{(k)} - \\alpha p^{(k)})\n",
    "$$\n",
    "leading again to the choice of $\\alpha$ of\n",
    "$$\n",
    "    \\alpha^{(k)} = \\frac{(p^{(k)})^T p^{(k)}}{(p^{(k)})^T A p^{(k)}}\n",
    "$$\n",
    "accept now we are also allowed to pick the search direction $p^{(k)}$.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Ways to choose $p^{(k)}$:\n",
    " - One bad choice for $p$ would be orthogonal to $r$ since this would then be tangent to the level set (ellipse) of $\\phi(u^{(k)})$ and would cause it to only increase so we want to make sure that $p^T r \\neq 0$ (the inner-product).\n",
    " - We also want to still move downwards so require that $\\phi(u^{(k+1)}) < \\phi(u^{(k)})$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We know that $r^{(k)}$ is not always the best direction to go in but what might be better?  \n",
    "\n",
    "We could head directly for the minimum but how do we do that without first knowing $u^\\ast$?  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It turns out when $m=2$ we can do this an from any initial guess $u^{(0)}$ and initial direction $p^{(0)}$ we will arrive at the minimum in 2 steps if we choose the second search direction dependent on\n",
    "$$\n",
    "    (p^{(1)})^T A p^{(0)} = 0.\n",
    "$$\n",
    "In general if two vectors satisfy this property they are said to be $A$-conjugate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Note that if $A = I$ then these two vectors would be orthogonal to each other and in this sense $A$-conjugacy is a natural extension from orthogonality and the simple case from before where the ellipses are circles to the case where we can have very distorted ellipses.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In fact the vector $p^{(0)}$ is tangent to the level set that $u^{(1)}$ lies on and therefore choosing $p^{(1)}$ so that it is $A$-conjugate to $p^{(0)}$ then always heads to the center of the ellipse. \n",
    "\n",
    "To show this, take the initial direction $p^{(0)}$ and note that it is tangent to the level set of $\\phi$ at $u_1$ so $p^{(0)}$ is orthogonal to the residual.  Looking at the residual of one step we have\n",
    "$$\n",
    "    r_1 = f - A u_1 = Au^\\ast - A u_1 = A(u^\\ast - u_1) \\Rightarrow (p^{(0)})^T A (u^\\ast - u_1) = 0.\n",
    "$$\n",
    "\n",
    "Knowing then that $u^\\ast - u_1 = \\alpha p_1$ we can then conclude\n",
    "$$\n",
    "    p_0^T A p_1 = 0\n",
    "$$\n",
    "given that $\\alpha \\neq 0$.\n",
    "\n",
    "The fact that the initial guess minimizes the residual in its direction is an important idea we will see again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In other words, once we know a tangent to one of the ellipses we can always choose a direction that minimizes in one of the dimensions of the search space.  Choosing the $p^{(k)}$ iteratively this way forms the basis of the **conjugate gradient** method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "![Ellipses CG](./images/ellipses_CG.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now to generalize beyond $m = 2$ consider the $m=3$ case.  As stated before we are now in a three-dimensional space where the level-sets are concentric ellipsoids.  Taking a slice through this space will lead to an ellipse on the slice.\n",
    "\n",
    "1. Start with an initial guess $u^{(0)}$ and choose a search direction $p^{(0)}$.\n",
    "1. Minimize $\\phi(u)$ in the direction $u^{(0)} + \\alpha p^{(0)}$ resulting in the choice\n",
    "$$\n",
    "    \\alpha^{(0)} = \\frac{(p^{(0)})^T p^{(0)}}{(p^{(0)})^T A p^{(0)}}\n",
    "$$\n",
    "as we saw before.  Now set $u^{(1)} = u^{(0)} + \\alpha^{(0)} p^{(0)}$.\n",
    "1. Choose $p^{(1)}$ to be $A$-conjugate to $p^{(0)}$.  In this case there are an infinite set of vectors that are possible that satisfy $(p^{(1)})^T A p^{(0)} = 0$.  Beyond requiring $p^{(1)}$ to be $A$-conjugate we also want it to be linearly-independent to $p^{(0)}$.\n",
    "1. Again choose an $\\alpha^{(1)}$ that minimizes the residual (again tangent to the level sets of $\\phi$) in the direction $p^{(1)}$ and repeat the process"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "So why does this work so much better than steepest descent?\n",
    "\n",
    "Since we have chosen $p^{(0)}$ and $p^{(1)}$ to be linearly independent they span a two-dimensional subspace in $\\mathbb R^m$.  We then choose to travel through this subspace as to minimize the residual in this space.  Geometrically we have cut out space with this plane leading to a two-dimensional subspace with concentric ellipses representing the level sets of $\\phi$.  We then choose an $\\alpha$ so that we arrive at the center of these ellipses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "alpha = 0.0\n",
    "beta = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 50\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Construct matrix A\n",
    "A = numpy.zeros((m, m))\n",
    "diagonal = numpy.ones(m) / delta_x**2\n",
    "A += numpy.diag(diagonal * 2.0, 0)\n",
    "A += numpy.diag(-diagonal[:-1], 1)\n",
    "A += numpy.diag(-diagonal[:-1], -1)\n",
    "\n",
    "# Construct right hand side\n",
    "b = -f(x)\n",
    "b[0] += alpha / delta_x**2\n",
    "b[-1] += beta / delta_x**2\n",
    "\n",
    "# Exact solution to system\n",
    "U_true = numpy.linalg.solve(A, b)\n",
    "\n",
    "# Algorithm parameters\n",
    "MAX_ITERATIONS = 10000\n",
    "tolerance = delta_x**2\n",
    "\n",
    "# Solve system\n",
    "U = numpy.zeros(m)\n",
    "convergence_CG = numpy.zeros(MAX_ITERATIONS)\n",
    "step_size_CG = numpy.zeros(MAX_ITERATIONS)\n",
    "residual_norm = numpy.zeros(MAX_ITERATIONS)\n",
    "system_convergence = numpy.zeros(MAX_ITERATIONS)\n",
    "success = False\n",
    "r = -(b - numpy.dot(A, U))\n",
    "p = -r\n",
    "r_dot_r = numpy.dot(r, r)\n",
    "for k in range(MAX_ITERATIONS):\n",
    "    # Convergence measures (not needed)\n",
    "    U_old = U.copy()\n",
    "    r_old = r.copy()\n",
    "    \n",
    "    if numpy.linalg.norm(r, ord=2) < tolerance:\n",
    "        success = True\n",
    "        break\n",
    "    \n",
    "    A_dot_p = numpy.dot(A, p)\n",
    "    \n",
    "    # Step length\n",
    "    alpha =  r_dot_r / numpy.dot(p, A_dot_p)\n",
    "    \n",
    "    # Update solution and residual\n",
    "    U += alpha * p\n",
    "    r += alpha * A_dot_p\n",
    "\n",
    "    # Improvement of solution\n",
    "    beta = numpy.dot(r, r) / r_dot_r\n",
    "    r_dot_r = numpy.dot(r, r)\n",
    "    \n",
    "    # Pick new direction for next iteration (if needed)\n",
    "    p = beta * p - r\n",
    "\n",
    "    # Convergence measures\n",
    "    convergence_CG[k] = numpy.linalg.norm(u_true(x) - U, ord=2)\n",
    "    residual_norm[k] = numpy.linalg.norm(r, ord=2)\n",
    "    step_size_CG[k] = numpy.linalg.norm(U - U_old, ord=2)\n",
    "    system_convergence[k] = numpy.linalg.norm(U_true - U, ord=2)\n",
    "\n",
    "if not success:\n",
    "    print(\"Iteration failed to converge!\")\n",
    "    print(residual_norm[-10:])\n",
    "    print(convergence_CG[-10:])\n",
    "else:\n",
    "    # Plot result\n",
    "    fig = plt.figure()\n",
    "    axes = fig.add_subplot(1, 1, 1)\n",
    "    axes.plot(x, U, 'o', label=\"Computed\")\n",
    "    axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "    axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "    axes.set_xlabel(\"x\")\n",
    "    axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "    fig = plt.figure()\n",
    "    fig.set_figwidth(fig.get_figwidth() * 2.1)\n",
    "    fig.set_figheight(fig.get_figheight() * 2.2)\n",
    "    \n",
    "    axes = fig.add_subplot(2, 2, 1)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), step_size_CG[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Step Size\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "    \n",
    "    axes = fig.add_subplot(2, 2, 2)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), convergence_CG[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Convergence to True Solution\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "    axes = fig.add_subplot(2, 2, 3)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), residual_norm[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Residual\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||r||_2$\")\n",
    "\n",
    "    axes = fig.add_subplot(2, 2, 4)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), system_convergence[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"True Residual\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^* - U||_2$\")\n",
    "    \n",
    "    print(\"Matrix condition number = %s\" % (numpy.sqrt(numpy.linalg.norm(A) * numpy.linalg.norm(numpy.linalg.inv(A)))))\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def CG(A, U, b, axes):\n",
    "    MAX_ITERATIONS = A.shape[0] * 2\n",
    "    tolerance = 1e-8\n",
    "    success = False\n",
    "    iteration_locations = []\n",
    "    \n",
    "    r = -(b - numpy.dot(A, U))\n",
    "    p = -r\n",
    "    r_dot_r = numpy.dot(r, r)\n",
    "    for k in range(MAX_ITERATIONS):\n",
    "        axes.text(U[0] + 0.1, U[1] + 0.1, str(k), fontsize=12)\n",
    "        axes.plot(U[0], U[1], 'ro')\n",
    "        iteration_locations.append(U)\n",
    "        if numpy.linalg.norm(r, ord=2) < tolerance:\n",
    "            success = True\n",
    "            break\n",
    "\n",
    "        A_dot_p = numpy.dot(A, p)\n",
    "\n",
    "        # Step length\n",
    "        alpha =  r_dot_r / numpy.dot(p, A_dot_p)\n",
    "\n",
    "        # Update solution and residual\n",
    "        U += alpha * p\n",
    "        r += alpha * A_dot_p\n",
    "\n",
    "        # Improvement of solution\n",
    "        beta = numpy.dot(r, r) / r_dot_r\n",
    "        r_dot_r = numpy.dot(r, r)\n",
    "\n",
    "        # Pick new direction for next iteration (if needed)\n",
    "        p = beta * p - r\n",
    "    \n",
    "    if success:\n",
    "        return k, iteration_locations\n",
    "    else:\n",
    "        raise Exception(\"Iteration did not converge.\")\n",
    "\n",
    "phi = lambda X, Y, A: 0.5 * (A[0, 0] * X**2 + A[0, 1] * X * Y + A[1, 0] * X * Y + A[1, 1] * Y**2) - X * f[0] - Y * f[1]\n",
    "\n",
    "x = numpy.linspace(-15, 15, 100)\n",
    "y = numpy.linspace(-15, 15, 100)\n",
    "X, Y = numpy.meshgrid(x, y)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.identity(2)\n",
    "f = numpy.array([0.0, 0.0])\n",
    "k, iteration_locations = CG(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.array([[2, -1], [-1, 2]])\n",
    "f = numpy.array([1.0, 2.0])\n",
    "k, iteration_locations = CG(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "fig.set_figheight(fig.get_figheight() * 2)\n",
    "axes = fig.add_subplot(1, 1, 1, aspect='equal')\n",
    "A = numpy.array([[2, -1.8], [-1.8, 2]])\n",
    "f = numpy.array([1.0, 2.0])\n",
    "k, iteration_locations = CG(A, numpy.array([-10.0, -12.0]), f, axes)\n",
    "axes.contour(X, Y, phi(X, Y, A), 25)\n",
    "print(\"Iteration count: %s\" % k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The conjugate gradient algorithm requires $\\mathcal{O}(m^2)$ operations and in theory $m$ iterates to achieve the exact solution.  In practice however it generally can converge much faster, especially when using preconditioning.  This is especially true for the two-dimensional, 5-point Laplacian stencil we have seen already."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Krylov Subspaces\n",
    "\n",
    "We can say a lot about the set of vectors $p^{(k)}$ that span the space our level set $\\phi$ lies in starting with the following theorem.\n",
    "\n",
    "**Theorem** The vectors generated in the CG algorithm have the following properties, provided $r^{(k)} \\neq 0$:\n",
    "1. $p^{(k)}$ is $A$-conjugate to all the previous search directions\n",
    "1. The residual $r^{(k)}$ is orthogonal to all the previous residuals $(r^{(k)})^T r_j = 0 \\quad \\forall j=0,1,\\ldots$\n",
    "1. The following three subspaces of $\\mathbb R^m$ are identical\n",
    "$$\\begin{aligned}\n",
    "    &\\text{span} (p^{(0)}, p^{(1)}, \\ldots, p^{(k-1)}) \\\\\n",
    "    &\\text{span} (r^{(0)}, A r^{(0)}, A^2 r^{(0)}, \\ldots, A^{k-1} r^{(0)}) \\\\\n",
    "    &\\text{span} (A e^{(0)}, A^2 e^{(0)}, A^3 e^{(0)}, \\ldots, A^{k} e^{(0)})\n",
    "\\end{aligned}$$\n",
    "\n",
    "The subspace $\\mathcal{K}_k = \\text{span} (r^{(0)}, A r^{(0)}, A^2 r^{(0)}, \\ldots, A^{k-1} r^{(0)})$ spanned by the vector $r^{(0)}$ and the first $k-1$ powers of $A$ to $r^{(0)}$ is called a *Krylov space* of dimension $k$ associated with $r^{(0)}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Projections\n",
    "\n",
    "The new update $u^{(k)}$ is formed by adding multiples of the $p^{(j)}$ to the initial guess $u^{(0)}$ and therefore lies in the subspace $u_0 + \\mathcal{K}_k$.  Another way to write this would be to say $u^{(k)} - u^{(0)} \\in \\mathcal{K}_k$.  Note the similarity between this and our projection discussion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Convergence of Conjugate Gradient\n",
    "\n",
    "We now turn to the convergence of CG and to deriving estimates about the size of the error at a given $k$ and the rate of convergence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "First define the $A$-norm s.t.\n",
    "$$\n",
    "    ||e||_A = \\sqrt{e^T A e}\n",
    "$$\n",
    "relying on the fact that $A$ is SPD."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This leads to\n",
    "$$\\begin{aligned}\n",
    "    ||e||^2_A &= (u - u^\\ast)^T A (u - u^\\ast) \\\\\n",
    "    &=u^T A u - 2 u^T A u^\\ast + (u^\\ast)^T A u^\\ast \\\\\n",
    "    &= 2 \\phi(u) + (u^\\ast)^T A u^\\ast\n",
    "\\end{aligned}$$\n",
    "\n",
    "Given that $(u^\\ast)^T A u^\\ast$ is a constant minimizing the error $||e||_A$ is equivalent to minimizing $\\phi(u)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We know we can expand $u^{(k)}$ in the subpsace $u_0 + \\mathcal{K}_k$ so that\n",
    "$$\n",
    "    u^{(k)} = u^{(0)} + \\alpha^{(0)} p^{(0)} + \\alpha^{(1)} p^{(1)} + \\cdots + \\alpha^{(k-1)} p^{(k-1)}\n",
    "$$\n",
    "and subtracting $u^\\ast$ we find\n",
    "$$\n",
    "    u^{(k)} - u^\\ast = e^{(k)} = e^{(0)} + \\alpha^{(0)} p^{(0)} + \\alpha^{(1)} p^{(1)} + \\cdots + \\alpha^{(k-1)} p^{(k-1)}\n",
    "$$\n",
    "which relates $e^{(k)}$ and $e^{(0)}$ and implies that $e^{(k)} - e^{(0)} \\in \\mathcal{K}_k$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "By our theorem above we also know that\n",
    "$$\n",
    "    e^{(k)} - e^{(0)} \\in \\text{span} (A e^{(0)}, A^2 e^{(0)}, A^3 e^{(0)}, \\ldots, A^{k} e^{(0)})\n",
    "$$\n",
    "so that\n",
    "$$\n",
    "    e^{(k)} = e^{(0)} + c_1 A e^{(0)} + c_2 A^2 e^{(0)} + \\cdots + c_k A^k e^{(0)}\n",
    "$$\n",
    "where $c_j$ are some coefficients (found by taking the projection of $e^{(k)}$ onto the span)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This then leads to\n",
    "$$\n",
    "    e^{(k)} = P_k(A) e^{(0)}\n",
    "$$\n",
    "where\n",
    "$$\n",
    "    P_k(A) = I + c_1 A + c_2 A^2 + \\cdots + c_k A^k.\n",
    "$$\n",
    "Note that $P_k(A)$ is a polynomial in $A$.  We can relate this to a scalar polynomial by letting\n",
    "$$\n",
    "    P_k(x) = 1 + c_1 x + c_2 x^2 + \\cdots +c_k x^k\n",
    "$$\n",
    "so that $P_k \\in \\mathcal{P}_k$ where $\\mathcal{P}_k$ are polynomials of degree at most $k$ and $P(0) = 1$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "CG implicitly constructs the polynomial $P_k$ and solves the related minimization problem\n",
    "$$\n",
    "    \\min_{P \\in \\mathcal{P}_k} ||P(A) e^{(0)}||_A.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### What is a polynomial of matrices?\n",
    "\n",
    "If we know our matrix is diagonalizable we can write this diagonalization as\n",
    "$$\n",
    "    A = V \\Lambda V^{-1}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If we take powers of $A$ then we find\n",
    "$$\n",
    "    A^k = (V \\Lambda V^{-1})^k = V \\Lambda V^{-1} V \\Lambda V^{-1} \\cdots V \\Lambda V^{-1} = V \\Lambda^k V^{-1}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Polynomials in $A$ then can be written as\n",
    "$$\n",
    "    P_k(A) = V P_k(\\Lambda) V^{-1}\n",
    "$$\n",
    "so that\n",
    "$$\n",
    "    P_k(\\Lambda) = \\begin{bmatrix}\n",
    "        P_k(\\lambda_1) \\\\\n",
    "        & P_k(\\lambda_2) \\\\\n",
    "        & & P_k(\\lambda_3) \\\\\n",
    "        & & & \\ddots \\\\\n",
    "        & & & & P_k(\\lambda_m)\n",
    "    \\end{bmatrix}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now going back to our expression for the error we have\n",
    "$$\n",
    "    e^{(k)} = P_k(A) e^{(0)}.\n",
    "$$\n",
    "If the polynomial $P_k(A)$ has a root at each $\\lambda_k$ then $P_k(\\Lambda)$ is the zero matrix and the above expression for the error leads us to the conclusion that $e^{(k)} = 0$.  If the eigenvalues are repeated (say we have $n$ unique eigenvalues, this is known as geometric multiplicity) then we can say that there is again a polynomial $P_n \\in \\mathcal{P}_n$ that also has these roots and we would expect $e^{(n)} = 0$, we converge faster than the dimension of $A$!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Say we want to know how big the error is before the iteration that guarantees convergence?  We want to then know how $||e^{(0)}||_A$ behaves.  It turns out for any $P \\in \\mathcal{P}$ we have\n",
    "$$\n",
    "    \\frac{||P(A)e^{(0)}||_A}{||e^{(0)}||_A} \\leq \\max_{1 \\leq j \\leq m} |P(\\lambda_j)|.\n",
    "$$\n",
    "We then need to find one polynomial $\\hat{P_k} \\in \\mathcal{P}_k$ that we can obtain a useful bound on $||e^{(k)}||_A / ||e^{(0)}||_A$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Surprisingly enough shifted and scaled version of the Chebyshev polynomials $T_k(x)$ will work for this!  Setting\n",
    "$$\n",
    "    \\hat{P_k}(x) = \\frac{T_k\\left(\\frac{\\lambda_m + \\lambda_1 - 2x}{\\lambda_m - \\lambda_1} \\right )}{T_k\\left(\\frac{\\lambda_m + \\lambda_1}{\\lambda_m - \\lambda_1}\\right ) }\n",
    "$$\n",
    "\n",
    "and compute \n",
    "$$\n",
    "    \\max_{1\\leq j \\leq m} \\left |\\hat{P_k}(\\lambda_j) \\right|.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "With some analysis (refer to LeVeque for a full derivation) we can show that \n",
    "$$\n",
    "    T_k(x) = \\frac{1}{2} \\left[ \\left ( \\frac{\\sqrt{\\kappa} + 1}{\\sqrt{\\kappa} - 1} \\right)^k + \\left ( \\frac{\\sqrt{\\kappa} - 1}{\\sqrt{\\kappa} + 1} \\right)^k \\right ]\n",
    "$$\n",
    "and\n",
    "$$\n",
    "    \\frac{||P(A) e^{(0)}||_A}{||e^{(0)}||_A} \\leq 2 \\left[ \\left ( \\frac{\\sqrt{\\kappa} + 1}{\\sqrt{\\kappa} - 1} \\right)^k + \\left ( \\frac{\\sqrt{\\kappa} - 1}{\\sqrt{\\kappa} + 1} \\right)^k \\right ]^{-1} \\leq 2 \\left ( \\frac{\\sqrt{\\kappa} - 1}{\\sqrt{\\kappa} + 1} \\right)^k\n",
    "$$\n",
    "where $\\kappa$ is the matrix condition number."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If the condition number of the matrix $\\kappa$ is large we can also simplify our bound to\n",
    "$$\n",
    "    2 \\left ( \\frac{\\sqrt{\\kappa} - 1}{\\sqrt{\\kappa} + 1} \\right)^k \\approx 2 \\left ( 1 - \\frac{2}{\\sqrt{\\kappa}} \\right)^k \\approx 2 e^{-2 k / \\sqrt{\\kappa}}.\n",
    "$$\n",
    "This implies that the expected number of iterations $k$ to reach a desired tolerance will be $k = \\mathcal{O}\\left(\\sqrt{\\kappa}\\right)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Preconditioning\n",
    "\n",
    "One way to get around the difficulties with these types of methods due to the distortion of the ellipses (and consequently the conditioning of the matrix) is to precondition the matrix.  The basic idea is that we take our original problem $A u = f$ and instead solve\n",
    "$$\n",
    "    M^{-1} A u = M^{-1} f.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Note that since we need to find the inverse of $M$, this matrix should be nice.  A couple of illustrative examples may help to illustrate why this might be a good idea:\n",
    "\n",
    " - If $M = A$ then we essentially have solved our problem already although that does not help us much\n",
    " - If $M = \\text{diag}(A)$, then $M^{-1}$ is easily computed and it turns out for some problems this can decrease the condition number of $M^{-1} A$ significantly.  Note though that this is not actually helpful in the case of the Poisson problem.\n",
    " - If $M$ is based on another iterative method used on $A$, for instance Gauss-Seidel, these can be effective general use preconditioners for many problems.\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "So the next question then becomes how to choose a preconditioner.  This is usually very problem specific and a number of papers suggest strategies for particular problems.  In general we want to move the eigenvalues around so that the matrix is not as ill-conditioned.  Combining a preconditioner with CG leads to the popular PCG method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Generalized Minimum Residual (GMRES) Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What if our system is not symmetric positive (negative) definite?\n",
    "\n",
    "Panic! \n",
    "\n",
    "(or don't if you have a towel handy and have a particular guide handy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Not really.  GMRES is one approach to solving this problem for us.  Since we can no longer depend on the structure of the related scalar minimization problem we instead will minimize the residual in the same subspaces we had before using a least-squares approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In the $k$th iteration GMRES solves a least squares problem in a particular subspace of the full problem.  In this case we will use the spaces we saw before, $u_0 + \\mathcal{\\kappa}_k$, where $\\mathcal{\\kappa}$ is defined as\n",
    "$$\n",
    "    \\kappa_k = \\text{span} (r^{(0)}, A r^{(0)}, A^2 r^{(0)}, \\ldots, A^{k-1} r^{(0)}).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To do this we construct a matrix $Q$ that contains in its columns a set of orthonormal vectors that span the space $\\kappa_k$.  Since this is iterative each time we increment $k$ we only need to find one additional vector for $Q$ to span the next space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Starting with the $k$th iterate we have\n",
    "$$\n",
    "    Q = [q_1 q_2 q_3 \\cdots~q_k] \\in \\mathbb R^{m \\times k}.\n",
    "$$\n",
    "Pick a vector $v_j \\notin \\kappa_k$ and use Gram-Shmidt (or something similar) to orthogonalize this vector to all the rest of the $q_i$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What should we choose for $v_j$?\n",
    "\n",
    "There are some bad choices such as $v_k = A^k r^{(k)}$ (this works but leads to very small singular values).\n",
    "\n",
    "Instead we will choose $v_k = A q_k$.  In effect we are building up a factorization of the matrix $A$.  This process is called the *Arnoldi process*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Arnoldi Process\n",
    "\n",
    "The Arnoldi process builds an orthogonal basis of the Krylov subspace $\\mathcal{K}_m$.  We can use any of the algorithms to find the orthogonalization (Gram-Schmidt for instance).  In essence then we have the following procedure\n",
    "\n",
    "Initialize process with a vector $v_0$ with $||v_0|| = 1$.  Now iterate on the following steps:\n",
    "1. $w_k = A v_{k-1}$\n",
    "2. Orthogonalize the vector $w_j$ against all the previous $v_j \\quad j < k$\n",
    "3. Stop when $w_j$ vanishes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This process has the practical effect of creating the following transformation\n",
    "$$\n",
    "    A V_m = V_m H_m\n",
    "$$\n",
    "where $V_m$ is the $m \\times m$ matrix with the columns the computed vectors $v_j$ and $H_m$ is the $(m + 1) \\times m$ Hessenberg matrix whose nonzero entries are the $h_{ij}$ defined as $h_{ij} = A v_j \\cdot v_i$ and $h_{j+1,j} = ||w_j||_2$ with the deletion of the last row.  This describes exactly an orthogonalization process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now how do we apply this?  Given an initial guess $x^{(0)}$ consider the orthogonal projection method where\n",
    "$$\n",
    "    \\mathcal{L} = \\mathcal{K} = \\mathcal{K}_m(A, r^{(0)})\n",
    "$$\n",
    "where the Krylov subspace is\n",
    "$$\n",
    "    \\mathcal{K}_m(A, r^{(0)}) = \\text{span}\\{r^{(0)}, A r^{(0)}, A^2 r^{(0)}, \\ldots,A^{m-1} r^{(0)} \\}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Let $v_1 = r^{(0)} / || r^{(0)} ||_2$ and apply Arnoldi's method we then would have\n",
    "$$\n",
    "    V^T_m A V_m = H_m\n",
    "$$\n",
    "which also implies that\n",
    "$$\n",
    "    V^T_m r^{(0)} = V^T_m (\\beta v_1) = \\beta e_1\n",
    "$$\n",
    "where $\\beta = ||r^{(0)}||_2$.  We can then find\n",
    "$$\n",
    "    x_m = x^{(0)} + V_m y_m \\\\\n",
    "    y_m = H^{-1}_m (\\beta e_1)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This algorithm is equivalent to Conjugate Gradient when the matrix is SPD.  \n",
    "\n",
    "Also note that in the case the matrix is symmetric the Hessenberg matrix $H$ is also symmetric implying that it is tri-diagonal greatly simplifying our general result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def arnoldi(A, b):\n",
    "    n = b.shape[0]\n",
    "    q = [0] * n\n",
    "    q[0] = b / numpy.linalg.norm(b)\n",
    "    h = numpy.zeros((n + 1, n))\n",
    "    for k in range(n):\n",
    "        v = numpy.dot(q[k], A).reshape(-1)\n",
    "        for j in range(k):\n",
    "            h[j, k] = numpy.dot(q[j], v.T)\n",
    "            v = v - h[j, k] * q[j]      \n",
    "        h[k + 1, k] = numpy.linalg.norm(v)\n",
    "        if (h[k + 1, k] != 0 and k != n - 1):\n",
    "            q[k + 1] = v / h[k + 1, k]\n",
    "    \n",
    "    return h\n",
    "\n",
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "alpha = 0.0\n",
    "beta = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 100\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Construct matrix A\n",
    "A = numpy.zeros((m, m))\n",
    "diagonal = numpy.ones(m) / delta_x**2\n",
    "A += numpy.diag(diagonal * 2.0, 0)\n",
    "A += numpy.diag(-diagonal[:-1], 1)\n",
    "A += numpy.diag(-diagonal[:-1], -1)\n",
    "\n",
    "# Construct right hand side\n",
    "b = -f(x)\n",
    "b[0] += alpha / delta_x**2\n",
    "b[-1] += beta / delta_x**2\n",
    "\n",
    "U = numpy.zeros(m + 2)\n",
    "U[0] = alpha\n",
    "U[-1] = beta\n",
    "\n",
    "# Compute Hessenberg matrix\n",
    "H = arnoldi(A, b)\n",
    "\n",
    "# Solve resulting minimization problem\n",
    "e_1 = numpy.zeros(m)\n",
    "e_1[0] = 1.0\n",
    "y = numpy.linalg.solve(H, beta * e_1)\n",
    "# NOT WORKING!\n",
    "\n",
    "# V = numpy.zeros((m, m))\n",
    "# H = numpy.zeros((m, m))\n",
    "\n",
    "# tolerance = 1e-14\n",
    "# r = b - numpy.dot(A, U[1:-1])\n",
    "# beta = numpy.linalg.norm(r, ord=2)\n",
    "# V[:, 0] = r / beta\n",
    "# for j in range(m):\n",
    "#     # Krylov subspace\n",
    "#     w = numpy.dot(A, V[:, j])\n",
    "#     # Modified Gram-Schmidt\n",
    "#     for i in range(j + 1):\n",
    "#         H[i, j] = numpy.dot(w, V[:, i])\n",
    "#         w = w - H[i, j] * V[:, i]\n",
    "#     # Need to handle some indexing pieces\n",
    "#     if j == m - 1:\n",
    "#         break\n",
    "#     H[j + 1, j] = numpy.linalg.norm(w, ord=2)\n",
    "#     # If this happens then the orthogonalization has completed and\n",
    "#     # something else may need to be done\n",
    "#     if numpy.abs(H[j + 1, j]) < tolerance:\n",
    "#         break\n",
    "#     # Have next basis vector\n",
    "#     V[:, j + 1] = w / H[j + 1, j]\n",
    "    \n",
    "# # Solve resulting minimization problem\n",
    "# e_1 = numpy.zeros(m)\n",
    "# e_1[0] = 1.0\n",
    "# y = numpy.linalg.solve(H, beta * e_1)\n",
    "# U[1:-1] += numpy.dot(V, y)\n",
    "\n",
    "# Plot result\n",
    "fig = plt.figure()\n",
    "axes = fig.add_subplot(1, 1, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Multigrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Jacobi Revisited\n",
    "\n",
    "Consider the Poisson problem with\n",
    "$$\n",
    "    f(x) = -20 + \\frac{1}{2}\\left ( \\phi''(x) \\cos \\phi(x) -  (\\phi'(x))^2 \\sin \\phi(x) \\right )\n",
    "$$\n",
    "with\n",
    "$$\n",
    "    \\phi = 20 \\pi x^3,\n",
    "$$\n",
    "and boundary conditions $u(0) = 1$ and $u(1) = 3$.\n",
    "\n",
    "Integrating twice we can find that the solution of this problem is\n",
    "$$\n",
    "    u(x) = 1 + 12 x - 10 x^3 + \\frac{1}{2} \\sin \\phi(x).\n",
    "$$\n",
    "\n",
    "Discretizing this problem in the standard way with second order, centered finite differences leads to the following code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def jacobi_update(x, U, f, delta_x):\n",
    "    \"\"\"Update U with a single Jacobi iteration\"\"\"\n",
    "    U_new = U.copy()\n",
    "    for i in range(1, x.shape[0] - 1):\n",
    "        U_new[i] = 0.5 * (U[i+1] + U[i-1]) - f(x[i]) * delta_x**2 / 2.0\n",
    "    step_size = numpy.linalg.norm(U_new - U, ord=2)\n",
    "    del(U)\n",
    "    return U_new, step_size\n",
    "\n",
    "\n",
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "alpha = 1.0\n",
    "beta = 3.0\n",
    "phi = lambda x: 20.0 * numpy.pi * x**3\n",
    "phi_prime = lambda x: 60.0 * numpy.pi * x**2\n",
    "phi_dbl_prime = lambda x: 120.0 * numpy.pi * x\n",
    "f = lambda x: -20.0 + 0.5 * (phi_dbl_prime(x) * numpy.cos(phi(x)) - (phi_prime(x))**2 * numpy.sin(phi(x)))\n",
    "u_true = lambda x: 1.0 + 12.0 * x - 10.0 * x**2 + 0.5 * numpy.sin(phi(x))\n",
    "\n",
    "# Descretization\n",
    "m = 100\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Expected iterations needed\n",
    "iterations_J = int(2.0 * numpy.log(delta_x) / numpy.log(1.0 - 0.5 * numpy.pi**2 * delta_x**2))\n",
    "# iterations_J = 100\n",
    "\n",
    "# Solve system\n",
    "# Initial guess for iterations\n",
    "U = 1.0 + 2.0 * x_bc\n",
    "U[0] = alpha\n",
    "U[-1] = beta\n",
    "convergence_J = numpy.zeros(iterations_J)\n",
    "step_size = numpy.zeros(iterations_J)\n",
    "for k in range(iterations_J):\n",
    "    U, step_size[k] = jacobi_update(x_bc, U, f, delta_x)\n",
    "    convergence_J[k] = numpy.linalg.norm(delta_x * (u_true(x_bc) - U), ord=2)\n",
    "\n",
    "# Plot result\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 3)\n",
    "\n",
    "axes = fig.add_subplot(1, 3, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = f(x)$\")\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "axes = fig.add_subplot(1, 3, 2)\n",
    "axes.semilogy(list(range(iterations_J)), convergence_J, 'o')\n",
    "axes.set_title(\"Error\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - u(x)||_1$\")\n",
    "\n",
    "axes = fig.add_subplot(1, 3, 3)\n",
    "axes.semilogy(list(range(iterations_J)), step_size, 'o')\n",
    "axes.set_title(\"Change Each Step\")\n",
    "axes.set_xlabel(\"Iteration\")\n",
    "axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We eventually converge although we see that there is a lower limit to the effectiveness of the Jacobi iterations.  We also again observe the extremely slow convergence we expect.  What if we could still take advantage of Jacobi though?  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "alpha = 1.0\n",
    "beta = 3.0\n",
    "phi = lambda x: 20.0 * numpy.pi * x**3\n",
    "phi_prime = lambda x: 60.0 * numpy.pi * x**2\n",
    "phi_dbl_prime = lambda x: 120.0 * numpy.pi * x\n",
    "f = lambda x: -20.0 + 0.5 * (phi_dbl_prime(x) * numpy.cos(phi(x)) - (phi_prime(x))**2 * numpy.sin(phi(x)))\n",
    "u_true = lambda x: 1.0 + 12.0 * x - 10.0 * x**2 + 0.5 * numpy.sin(phi(x))\n",
    "\n",
    "# Descretization\n",
    "m = 255\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "U = 1.0 + 2.0 * x_bc\n",
    "U[0] = alpha\n",
    "U[-1] = beta\n",
    "\n",
    "num_steps = 1000\n",
    "plot_frequency = 200\n",
    "\n",
    "# Plot initial error\n",
    "fig = plt.figure()\n",
    "fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 1)\n",
    "axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "axes.set_title(\"Solution to $u_{xx} = f(x)$, iterations = %s\" % 0)\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "axes = fig.add_subplot(1, 2, 2)\n",
    "axes.plot(x_bc, U - u_true(x_bc), 'r-o')\n",
    "axes.set_title(\"Error, iterations = %s\" % 0)\n",
    "axes.set_xlabel(\"x\")\n",
    "axes.set_ylabel(\"U - u\")\n",
    "\n",
    "# Start Jacobi iterations\n",
    "for k in range(num_steps):\n",
    "    U, step_size = jacobi_update(x_bc, U, f, delta_x)\n",
    "\n",
    "    if (k+1)%plot_frequency == 0:\n",
    "        print(step_size)\n",
    "        fig = plt.figure()\n",
    "        fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "\n",
    "        axes = fig.add_subplot(1, 2, 1)\n",
    "        axes.plot(x_bc, U, 'o', label=\"Computed\")\n",
    "        axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "        axes.set_title(\"Solution to $u_{xx} = f(x)$, iterations = %s\" % (k + 1))\n",
    "        axes.set_xlabel(\"x\")\n",
    "        axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "        axes = fig.add_subplot(1, 2, 2)\n",
    "        axes.plot(x_bc, U - u_true(x_bc), 'r-o')\n",
    "        axes.set_title(\"Error, iterations = %s\" % (k + 1))\n",
    "        axes.set_xlabel(\"x\")\n",
    "        axes.set_ylabel(\"U - u\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Note that higher frequency components of the error are removed first!  Why might this be?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Recall that we found in general that the error $e^{(k)}$ from a matrix splitting iterative approach involves the matrix $G$ where\n",
    "$$\n",
    "    U^{(k+1)} = M^{-1} N U^{(k)} + M^{-1} b = G U^{(k)} + c.\n",
    "$$\n",
    "We then know that\n",
    "$$\n",
    "    e^{(k)} = G e^{(k-1)}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In the case for Jacobi the matrix $G$ can be written as\n",
    "$$\n",
    "    G = I + \\frac{\\Delta x^2}{2} A = \\begin{bmatrix}\n",
    "        0 & 1/2 & \\\\\n",
    "        1/2 & 0 & 1/2 \\\\\n",
    "         & 1/2 & 0 & 1/2 \\\\\n",
    "         & & \\ddots & \\ddots & \\ddots \\\\\n",
    "         & & & 1/2 & 0 & 1/2 \\\\\n",
    "         & & & & 1/2 & 0\n",
    "    \\end{bmatrix}\n",
    "$$\n",
    "Note that this amounts to averaging the off diagonal terms $U_{i+1}$ and $U_{i-1}$.  Averaging has the effect of smoothing, i.e. it damps out higher frequencies more quickly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Recall that the eigenvectors of $A$ and $G$ are the same, if those eigenvectors are\n",
    "$$\n",
    "    u^p_j = \\sin(\\pi p x_j) \\quad \\text{with} \\quad x_j = j \\Delta x, \\quad j = 1, 2, 3, \\ldots, m.\n",
    "$$\n",
    "with eigenvalues\n",
    "$$\n",
    "    \\lambda_p = \\cos(p \\pi \\Delta x).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can project the initial error $e^{(0)}$ onto the eigenspace such that\n",
    "$$\n",
    "    e^{(0)} = c_1 u^1 + c_2 u^2 + \\cdots + c_m u^m\n",
    "$$\n",
    "and therefore\n",
    "$$\n",
    "    e^{(k)} = c_1 (\\lambda_1)^k u^1 + c_2 (\\lambda_2)^k u^2 + \\cdots + c_m (\\lambda_m)^k u^m.\n",
    "$$\n",
    "This implies that the $p$th component of the vector $e^{(k)}$ decays as the corresponding eigenvalue."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Examining the eigenvalues we know that the 1st and $m$th eigenvalues will be closest to 1 so the terms $c_1 (\\lambda_1)^k u^1$ and $c_m (\\lambda_m)^k u^m$ will dominate the error as\n",
    "$$\n",
    "    \\lambda_1 = -\\lambda_m \\approx 1- \\frac{1}{2} \\pi^2 \\Delta x^2.\n",
    "$$\n",
    "We saw this before as this determined the overall convergence rate for Jacobi."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For other components of the error we can approximately see how fast they will decay.  If $m / 4 \\leq p \\leq 3m / 4$ then\n",
    "$$\n",
    "    |\\lambda_p| \\leq \\frac{1}{\\sqrt{2}} \\approx 0.7\n",
    "$$\n",
    "implying that after 20 iterations we would have $|\\lambda_p|^20 < 10^{-3}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Connecting this back our original supposition that higher order frequencies in the error are damped more quickly look at the form of the eigenvector components.  The original error was projected onto the eigenvectors so that\n",
    "$$\n",
    "    e^{(0)} = c_1 u^1 + c_2 u^2 + \\cdots + c_m u^m,\n",
    "$$\n",
    "plugging in the eigenvectors themselves we find\n",
    "$$\n",
    "    e^{(0)} = c_1 \\sin(\\pi x_j) + c_2 \\sin(\\pi 2 x_j) + \\cdots + c_m \\sin(\\pi m x_j)\n",
    "$$\n",
    "so that we have effectively broken down the original error in terms of a Fourier sine series.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now considering our analysis on the eigenvalues we see that it is in fact the middle range of frequencies that decay the most quickly.  The reason we did not see this in our example is that the solution did not contain high-order frequencies relative to our choice of $m$.  If we picked an $m$ that was too small we would have been in trouble (for multiple reasons).  Try this out and see what you observe."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It turns out that having only the middle ranges decay quickly is suboptimal in the context we are considering.  Instead we will user *underrelaxed Jacobi*, similar to SOR from before, where\n",
    "$$\n",
    "    U^{(k+1)} = (1 - \\omega) U^{(k)} + \\omega G U^{(k)}\n",
    "$$\n",
    "with $\\omega = 2/3$ (where $G$ is Jacobi's iteration matrix).  The new iteration matrix is\n",
    "$$\n",
    "    G_\\omega = (1 - \\omega) I + \\omega G\n",
    "$$\n",
    "and has eigenvalues\n",
    "$$\n",
    "    \\lambda_p = (1-\\omega)+\\omega \\cos (p \\pi \\Delta x).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This choice of $\\omega$ in fact then minimizes the eigenvalues in the range $m/2 < p < m$.  In fact\n",
    "$$\n",
    "    |\\lambda_p| \\leq 1/3\n",
    "$$\n",
    "for this range.  As a standalone method this is actually worse than Jacobi as the lower frequency components of the error decay even more slowly than before but this behavior is perfect for a multigrid approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Multigrid Approach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The basic approach for multigrid is to do the following:\n",
    "\n",
    "1. Use a few iterations of underrelaxed Jacobi to damp high-frequency components of the error.\n",
    "1. Since the error is now smoother than before we can represent the solution (and error) as a coarser resolution grid.  We then switch to a coarser resolution.\n",
    "1. On this coarser grid we then again apply a few iterations of underrelaxed Jacobi which now quickly removes a set of lower-frequency components of the error since we are on a coarser grid."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Consider the $p = m/4$ component that on the original grid will not be damped much since it is in the lower-frequency range.  If we transfer the problem now to a grid with half as many points we then find that the previous frequency is now at the midpoint of the frequency range and therefore will be damped much more quickly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "One important point here is that instead of transferring the solution to the coarser grid we only transfer the error.\n",
    "\n",
    "If we have taken $n$ iterations on the original grid we would now have\n",
    "$$\n",
    "    e^{(n)} = U^{(n)} - u\n",
    "$$\n",
    "which has a residual vector of\n",
    "$$\n",
    "    r^{(n)} = f - A U^{(n)}\n",
    "$$\n",
    "since\n",
    "$$\n",
    "    A e^{(n)} = -r^{(n)}.\n",
    "$$\n",
    "If we can solve this system for $e^{(n)}$ then we could go back and simply subtract this from the equation relating our numerical solution $U^{(n)}$ to $u$.  The system of equations relating the error and residual is the one we are interested in coarsening."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Basic Algorithm:\n",
    "1. Take $n$ iterations (where $n$ is fixed) of a simple iterative method on the original problem $A u = f$.  This gives the approximate solution $U^{(n)} \\in \\mathbb R^m$.\n",
    "1. Compute the residual $r^{(n)} = f - A U^{(n)} \\in \\mathbb R^m$.\n",
    "1. Coarsen the residual problem, take $r^{(n)} \\in \\mathbb R^m$ to $\\hat{r~} \\in \\mathbb R^{m_c}$ where $m_c = (m - 1) / 2$.\n",
    "1. Approximately solve the new problem $\\hat{A~} \\hat{e~} = -\\hat{r~}$ where $\\hat{A~}$ is the appropriately scaled matrix $A$.\n",
    "1. We now have an approximation to the error $e^{(k)}$ in $\\hat{e~}$.  To get back to $e^{(k)}$ we use an appropriate interpolation method to go back to $\\mathbb R^m$.  Now subtract this interpolated approximate error to get a new approximation $U$ to $u$.\n",
    "1. Using this new value of $U$ as an initial guess repeat the process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "There are many variations on this scheme, most notably that we do not need to stop at a single coarsening, we can continue to coarsen to dampen additional lower frequencies depending mostly on the number of grid points $m$ we started out with.  The specification of levels of coarsening combined with the interpolation back up to the original problem is called a *V-cycle*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Also note that although we are solving multiple linear systems, each coarsened system decreases the number of points used by half (this is also adjustable)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "![V-Cycle Example](./images/v-cycle.png)\n",
    "![W-Cycle Example](./images/w-cycle.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def smoother(v, N, A, f):\n",
    "    for i in range(N):\n",
    "        for k in range(A.shape[0]):\n",
    "            v[k] = (f[k] - numpy.dot(A[k,0:k], v[0:k]) - numpy.dot(A[k,k+1:], v[k+1:]) ) / A[k,k];        \n",
    "    return v\n",
    "\n",
    "def get_mat(A):\n",
    "    f = A.shape[0]\n",
    "    c =  int((f - 1) / 2 + 1)\n",
    "    P = numpy.zeros((f, c))\n",
    "    for k in range(c):\n",
    "        P[2 * k, k] = 1.0\n",
    "    for k in range(c - 1):\n",
    "        P[2 * k + 1, k] = 0.5\n",
    "        P[2 * k + 1, k + 1] = 0.5\n",
    "    return P  \n",
    "\n",
    "def vcycle(A, f):\n",
    "    if A.shape[0] < 15:\n",
    "        v = numpy.linalg.solve(A,f)\n",
    "        return v  \n",
    "    N = 5\n",
    "    v = numpy.zeros(A.shape[0]);\n",
    "    v = smoother(v, N, A, f)\n",
    "    P = get_mat(A)\n",
    "    remains = f - numpy.dot(A, v)\n",
    "    remains_P = numpy.dot(P.T, remains)\n",
    "    AP_P = numpy.dot(P.T, numpy.dot(A,P))\n",
    "    recurse_V = vcycle(AP_P, remains_P)\n",
    "    v = numpy.dot(P, recurse_V)\n",
    "    v = smoother(v, N, A, f)\n",
    "    return v       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "a = 0.0\n",
    "b = 1.0\n",
    "alpha = 0.0\n",
    "beta = 3.0\n",
    "f = lambda x: numpy.exp(x)\n",
    "u_true = lambda x: (4.0 - numpy.exp(1.0)) * x - 1.0 + numpy.exp(x)\n",
    "\n",
    "# Descretization\n",
    "m = 50\n",
    "x_bc = numpy.linspace(a, b, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "# Construct matrix A\n",
    "A = numpy.zeros((m, m))\n",
    "diagonal = numpy.ones(m) / delta_x**2\n",
    "A += numpy.diag(diagonal * 2.0, 0)\n",
    "A += numpy.diag(-diagonal[:-1], 1)\n",
    "A += numpy.diag(-diagonal[:-1], -1)\n",
    "\n",
    "# Construct right hand side\n",
    "b = -f(x)\n",
    "b[0] += alpha / delta_x**2\n",
    "b[-1] += beta / delta_x**2\n",
    "\n",
    "# True solution to system\n",
    "U_true = numpy.linalg.solve(A, b)\n",
    "\n",
    "# Algorithm parameters\n",
    "MAX_CYCLES = 50\n",
    "tolerance = delta_x**2\n",
    "\n",
    "# Perform multi-grid\n",
    "U = numpy.zeros(m)\n",
    "convergence_MG = numpy.zeros(MAX_CYCLES)\n",
    "step_size_MG = numpy.zeros(MAX_CYCLES)\n",
    "residual_norm = numpy.zeros(MAX_CYCLES)\n",
    "system_convergence = numpy.zeros(MAX_CYCLES)\n",
    "success = False\n",
    "for k in range(MAX_CYCLES):\n",
    "    U_old = U.copy()\n",
    "    \n",
    "    # Compute residual\n",
    "    r = b - numpy.dot(A, U)\n",
    "    if numpy.linalg.norm(r, ord=2) < tolerance:\n",
    "        success = True\n",
    "        break       \n",
    "    U += vcycle(A, r)\n",
    "    \n",
    "    # Convergence measures\n",
    "    convergence_MG[k] = numpy.linalg.norm(u_true(x) - U, ord=2)\n",
    "    residual_norm[k] = numpy.linalg.norm(r, ord=2)\n",
    "    step_size_MG[k] = numpy.linalg.norm(U - U_old, ord=2)\n",
    "    system_convergence[k] = numpy.linalg.norm(U_true - U, ord=2)\n",
    "\n",
    "if not success:\n",
    "    print(\"Iteration failed to converge!\")\n",
    "    print(residual_norm[-10:])\n",
    "    print(convergence_CG[-10:])\n",
    "else:\n",
    "    # Plot result\n",
    "    fig = plt.figure()\n",
    "    axes = fig.add_subplot(1, 1, 1)\n",
    "    axes.plot(x, U, 'o', label=\"Computed\")\n",
    "    axes.plot(x_bc, u_true(x_bc), 'k', label=\"True\")\n",
    "    axes.set_title(\"Solution to $u_{xx} = e^x$\")\n",
    "    axes.set_xlabel(\"x\")\n",
    "    axes.set_ylabel(\"u(x)\")\n",
    "\n",
    "    fig = plt.figure()\n",
    "    fig.set_figwidth(fig.get_figwidth() * 2.1)\n",
    "    fig.set_figheight(fig.get_figheight() * 2.2)\n",
    "    \n",
    "    axes = fig.add_subplot(2, 2, 1)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), step_size_MG[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Step Size\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "    \n",
    "    axes = fig.add_subplot(2, 2, 2)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), convergence_MG[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Convergence to True Solution\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^{(k)} - U^{(k-1)}||_2$\")\n",
    "\n",
    "    axes = fig.add_subplot(2, 2, 3)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), residual_norm[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"Residual\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||r||_2$\")\n",
    "\n",
    "    axes = fig.add_subplot(2, 2, 4)\n",
    "    axes.semilogy(numpy.arange(1, k + 1), system_convergence[:k], 'o')\n",
    "    axes.semilogy(numpy.arange(1, k + 1), numpy.ones(k) * delta_x**2, 'r--')\n",
    "    axes.set_title(\"True Residual\")\n",
    "    axes.set_xlabel(\"Iteration\")\n",
    "    axes.set_ylabel(\"$||U^* - U||_2$\")\n",
    "    \n",
    "    print(\"Matrix condition number = %s\" % (numpy.sqrt(numpy.linalg.norm(A) * numpy.linalg.norm(numpy.linalg.inv(A)))))\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Slightly harder problem\n",
    "L = 1.0\n",
    "m = 512\n",
    "x_bc = numpy.linspace(0, L, m + 2)\n",
    "x = x_bc[1:-1]\n",
    "delta_x = (b - a) / (m + 1)\n",
    "\n",
    "A = numpy.diag(2.0 * numpy.ones(m)) - numpy.diag(numpy.ones(m - 1),  1) / L**2\n",
    "\n",
    "b = numpy.tan(numpy.arange(m) / (2 * numpy.pi) * 10)\n",
    "\n",
    "U_true = numpy.linalg.solve(A, b)\n",
    "U = numpy.random.random(m)\n",
    "\n",
    "MAX_CYCLES = 50\n",
    "tolerance = 1e-10\n",
    "\n",
    "for k in range(MAX_CYCLES):\n",
    "    r = b - numpy.dot(A, U)\n",
    "    fig = plt.figure()\n",
    "    fig.set_figwidth(fig.get_figwidth() * 2)\n",
    "    axes = fig.add_subplot(1, 2, 1)\n",
    "    axes.plot(U)\n",
    "    axes.set_title(\"Solution at iteration %s\" % k)\n",
    "    axes = fig.add_subplot(1, 2, 2)\n",
    "    axes.plot(r)\n",
    "    axes.set_title(\"Residual at iteration %s\" % k)\n",
    "    \n",
    "    if numpy.linalg.norm(r) < tolerance:\n",
    "        break       \n",
    "\n",
    "    U += vcycle(A, r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Work Estimates\n",
    "\n",
    "So how much better is multigrid?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As a concrete example consider again the similar example from before where $m = 2^8 - 1 = 255$ and using $n = 3$.  If we allow recursive coarsening down to 3 grid points (7 levels of grids).  On each level we apply 3 iterations of Jacobi.  If we apply these Jacobi iterations on the way \"down\" the V-cycle and on the way up (not necessary in theory) we would do 6 iterations of Jacobi per level.  This leads to a total of 42 Jacobi iterations on variously coarsened grids.  The total number of updated values then would be\n",
    "$$\n",
    "    6 \\sum^8_{j=2} 2^j \\approx 3072.\n",
    "$$\n",
    "This is about the same amount of work as 12 iterations on the original fine grid would take.  The big difference here is that, due to the sweeping, we have in fact damped out error frequencies in a much larger range than would have been accomplished simply by Jacobi iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now consider the more general case with $m + 1 = 2^J$ points recursing all the way down to one point (about) and taking $n$ iterations of Jacobi at each level.  The total work would be\n",
    "$$\n",
    "    2 n \\sum^J_{j=2}2^j \\approx 4 n 2^J \\approx 4 n m = \\mathcal{O}(m)\n",
    "$$\n",
    "assuming $n \\ll m$.  The work required work for one V-cycle, noting that the number of grids grows as $\\log_2 m$, is still $\\mathcal{O}(m)$.  It can actually also be shown for the Poisson problem that the number of V-cycles required in our simple approach that $\\mathcal{O}(m \\log m)$ would be required to reach a given level of error determined by the original $\\Delta x$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can of course play with all sorts of types of V-cycles and iteration counts and these variations lead to a multitude of approaches. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Full Multigrid\n",
    "\n",
    "Instead of starting and solving the original PDE at the finest level we can also start at the coarsest.  To do this we do a few iterations on the coarsest level or solve the problem directly since the cost should be low at the coarsest level.  We then interpolate to the next finer level and solve the problem there.  We can then cycle back down to the coarsest level or continue upwards until we get to the finest level where we wanted to be in the first place.  We then switch back to solving for the error rather than the original problem and proceed as before.  This approach is usually labeled as *full multigrid* (FMG)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "![FMG-Cycle Example](./images/fmg-cycle.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It turns out that although there is a \"startup\" phase that this is mostly negligible and greatly reduces the error by the time we reach the finest level.  It turns out using FMG takes about $\\mathcal{O}(m)$ work, optimal given that we have $m$ unknowns to solve for in the 1-dimensional problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Take-Away"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We have been discussing one-dimensional implementations of multigrid but these methods extend to higher-dimensional problems and continue to be optimal.  For instance a two-dimensional Poisson problem can be solved in $\\mathcal{O}(m^2)$ work, again optimal due to the number of unknowns.  A Fourier transform approach would require $\\mathcal{O}(m^2 \\log m)$ and a direct method $\\mathcal{O}(m^3)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This all being said, multigrid is hard.  There are as many ways to do it as their are problems to be solved (at least).  Luckily there are a number of areas where research has been done to determine optimal methods (in some sense) for a given problem and discretization.  There are also variations that include more complex ways to \"coarsen\" the error and residual.  In general these are called *algebraic multigrid* methods (AMG).  These are especially useful when it is not obvious how to coarsen the residual problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Further Reading\n",
    "1. W. L. Briggs, V. Emden Henson, and S. F. McCormick. *A Multigrid Tutorial, 2nd ed.* SIAM, Philadelphia, 2000.\n",
    "1. D. C. Jesperson.  Multigrid methods for partial differential equations.  In *Studies in Numerical Analysis*, G.H. Golub, ed. MAA Studies in Mathematics, Vol. 24, 1984, pages 270-317.\n",
    "1. W. Hackbusch. *Multigrid Methods and Applications.* Springer-Verlag, Berlin, 1985.\n",
    "1. P. Wesseling. *An Introduction to Multigrid Methods*. John Wiley, New York, 1992."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.3"
  },
  "latex_envs": {
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 0
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}