{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using iterative methods to solve the Laplace equation\n",
    "\n",
    "## Theoretical introduction\n",
    "\n",
    "**The example in this notebook is based on the book [1]**.\n",
    "\n",
    "[1] Gabriele Ciaramella and Martin J. Gander. *Iterative Methods and Preconditioners for Systems of Linear Equations*. 2020."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will guide you through solving the **Laplace equation**. Consider the domain $\\Omega = (0,1)^2 \\subset \\mathbb{R}$ and a function $g \\in C^0(\\overline{\\Omega})$, we want to find a (strong) solution $u \\in C^2(\\overline{\\Omega})$ for the following boundary value problem\n",
    "\n",
    "\\begin{align*}\n",
    "\\Delta u &= 0 \\text{ in } \\Omega\\\\\n",
    "u &= g \\text{ on } \\partial \\Omega.\n",
    "\\end{align*}\n",
    "\n",
    "Here $\\Delta = \\partial_x^2 + \\partial_y^2$ is the **Laplace Operator**. As we are looking for a strong solution, we start with a finite difference approximation via\n",
    "\\begin{align*}\n",
    "\\partial_x^2 u(x,y) &\\approx \\frac{u(x+h, y) - 2u(x,y) + u(x-h,y)}{h^2} \\label{eq:DiscretePartialX} \\tag{1}\\\\\n",
    "\\partial_y^2 u(x,y) &\\approx \\frac{u(x, y+h) - 2u(x,y) + u(x,y-h)}{h^2}\\label{eq:DiscretePartialY} \\tag{2},\\\\\n",
    "\\end{align*}\n",
    "for a mesh size $h = \\frac{1}{m+1}> 0$, $m\\in \\mathbb{N}$. By defining the grid points $(x_j, y_i) = (jh, ih)$ and summing \\eqref{eq:DiscretePartialX} and \\eqref{eq:DiscretePartialY}, we obtain \n",
    "\\begin{align}\n",
    "u_{i,j-1} + u_{i-1,j} - 4u_{i,j} + u_{i+1,j} + u_{i,j+1}  = 0 \\label{eq:DiscreteLaplace}\\tag{3}\n",
    "\\end{align}\n",
    "\n",
    "as an approximation for equation $\\Delta u = 0$. For simplicity, we assume that $g \\equiv 1$ on the set $\\{(x,1) \\vert x \\in [0,1]\\}$ and $g \\equiv 0$ otherwise.\n",
    "\n",
    "Let us introduce the vectors\n",
    "\\begin{align*}\n",
    "\\mathbf{u}_j := \\begin{pmatrix}\n",
    "u_{1,j}\\\\\n",
    "\\vdots\\\\\n",
    "u_{m,j}\n",
    "\\end{pmatrix}, j = 1,2,\\ldots, m,\n",
    "\\end{align*}\n",
    "which contain the values $u_{i,j}$ for $i=1,\\ldots, m$ on an vertical column corresponding to $j$. By rewriting the upper system of linear equations \\eqref{eq:DiscreteLaplace} and using the top boundary condition (attention: **the boundary conditions are shifted into the right-hand side**), where $u \\equiv 1$, we get\n",
    "\\begin{align*}\n",
    "\\mathbf{u}_{j-1} + T \\mathbf{u}_j + \\mathbf{u}_{j+1} = - \\mathbf{e}_m = -(0,\\ldots, 0, 1)^T \\in \\mathbb{R}^m,\n",
    "\\end{align*}\n",
    "with the Matrix\n",
    "\\begin{align*}\n",
    "T := \\begin{pmatrix}\n",
    "-4 & 1 & & & \\\\\n",
    "1 & -4 & 1 & & \\\\\n",
    "& \\ddots & \\ddots & \\ddots & \\\\\n",
    "%& & \\ddots & \\ddots & 1 \\\\\n",
    "%& & & 1 & -4  \n",
    "\\end{pmatrix} \\in\\mathbb{R}^{m\\times m}.\n",
    "\\end{align*}\n",
    "By defining the vector\n",
    "\\begin{align*}\n",
    "\\mathbf{u} := \\begin{pmatrix}\n",
    "\\mathbf{u}_1\\\\\n",
    "\\vdots\\\\\n",
    "\\mathbf{u}_m\n",
    "\\end{pmatrix} \\in \\mathbb{R}^{m^2},\n",
    "\\end{align*}\n",
    "we obtain the following system of linear equations\n",
    "\\begin{align*}\n",
    "A \\mathbf{u} = \\mathbf{f}.\n",
    "\\end{align*}\n",
    "The matrix $A$ is given by\n",
    "\\begin{align*}\n",
    " A = \\text{I}_m \\otimes T + T \\otimes \\text{I}_m= \n",
    "\\begin{pmatrix}\n",
    "T & \\text{I}_m & & \\\\\n",
    "\\text{I}_m & T & \\text{I}_m & \\\\\n",
    "& \\ddots & \\ddots & \\ddots\n",
    "\\end{pmatrix} \\in\\mathbb{R}^{m^2\\times m^2},\n",
    "\\end{align*}\n",
    "with the **Kronecker product** $\\otimes$ and the right-hand side $\\mathbf{f}$ containing the boundary conditions. While constructing A, the boundedness of the domain has to be taken into account. In particular, that nodes at the edges do not have 4 neighbouring cells, but three or two."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us define the system matrices and the right-hand side. Feel free to change e.g. the mesh size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from oppy.tests.itMet import FDLaplacian\n",
    "\n",
    "# number of grid points\n",
    "m = 20\n",
    "# build the system matrix\n",
    "A = FDLaplacian(m, 2)\n",
    "# define the right-hand side\n",
    "f = np.zeros(m**2)\n",
    "f[m-1:-1:m] = -1\n",
    "# set the last entry manually\n",
    "f[-1] = -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute a first solution using the scipy sparse solver\n",
    "Fist, let us compute the exact solution using the *scipy sparse* solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enable 3D plotting\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib import cm\n",
    "\n",
    "from scipy.sparse.linalg import spsolve\n",
    "\n",
    "u_exact = spsolve(A, f)\n",
    "print('norm of the residual = {}'.format(np.linalg.norm(A.dot(u_exact)-f)))\n",
    "# we need to reshape the 1D vector back to 2D\n",
    "U_exact = np.zeros((m+2, m+2))\n",
    "U_exact[1:m+1, 1:m+1] = np.reshape(u_exact, (m, m))\n",
    "# manually set the correct boundary conditions\n",
    "U_exact[0:-1,-1] = 1\n",
    "\n",
    "x_grid = np.linspace(0,1,m+2)\n",
    "y_grid = x_grid\n",
    "# attention, later we need to interchange X,Y\n",
    "# (see the definition of the vector u)\n",
    "X, Y = np.meshgrid(x_grid, y_grid)\n",
    "\n",
    "\n",
    "# plotting\n",
    "fig = plt.figure(figsize=(10,10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlabel('x-axis')\n",
    "ax.set_ylabel('y-axis')\n",
    "ax.set_title('Exact solution of the Laplace equation using {} grid points'.format(m+1))\n",
    "surf = ax.plot_surface(Y,X, U_exact, cmap=cm.coolwarm)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iterative Solvers\n",
    "Now we want to use the solvers implemented in oppy to solve the upper system and compare their results. For the reason of comparsion, we run all of the algorithms with their default options and hence, we only define an options object once. We start with the standard iterative solver: Jacobi's method, Gauss-Seidel method, Successive over-relaxation (SOR) method and the Steepest descent algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from oppy.itMet import jacobi, gauss_seidel, sor, steepest_descent\n",
    "\n",
    "from oppy.options import Options\n",
    "\n",
    "# create a dict to save all results and print them in the end\n",
    "res_dict = {}\n",
    "\n",
    "jacobi_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, results='all')\n",
    "res_dict['jacobi'] = jacobi(A, f, jacobi_options)\n",
    "\n",
    "gauss_seidel_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, results='all')\n",
    "res_dict['gauss_seidel'] = gauss_seidel(A, f, gauss_seidel_options)\n",
    "\n",
    "sor_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, omega=1/2, results='all')\n",
    "res_dict['sor'] = sor(A, f, sor_options)\n",
    "\n",
    "steepest_descent_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, results='all')\n",
    "res_dict['steepest_descent'] = steepest_descent(A, f, steepest_descent_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Krylov methods\n",
    "In order to further improve our results, we want to switch from stationary solvers to more advanced krylov (subspace) methods: the conjugated gradient (CG) method and the generalized minimal residual (GMRES) algorithm. For the reason of comparability, we will keep the same options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from oppy.itMet import cg, gmres\n",
    "cg_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, results='all')\n",
    "res_dict['cg'] = cg(A, f, cg_options)\n",
    "\n",
    "gmres_options = Options(tol_abs=1e-7, tol_rel=1e-7, nmax=1e5, results='all')\n",
    "res_dict['gmres'] = gmres(A, f, gmres_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion: Comparison of convergence rates\n",
    "Let us plot the different convergence behaviours of our methods. In the plot below, you can see the history with the norm of the residuals and plotted against the iteration index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the norm of the residual for all methods:\n",
    "plt.figure(figsize=(15,10))\n",
    "\n",
    "for key in res_dict:\n",
    "    plt.semilogy(res_dict[key].res, label = key)\n",
    "    \n",
    "\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title('Convergence results for different methods')\n",
    "plt.xlabel('Number of iterations')\n",
    "plt.ylabel(\"Norm of the residual\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One can see that the Krylov methods provide really good convergence results and a convincing precision.\n",
    "## Krylov methods vs. scipy sparse solver\n",
    "In order to improve the results even further and compete with the ``scipy.sparse.linalg.spsolve`` method, we will increase the precision\n",
    "of our algorithms even more:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# again the scipy solution\n",
    "u_exact = spsolve(A, f)\n",
    "print('norm of the residual (scipy) = {}'.format(np.linalg.norm(A.dot(u_exact)-f)))\n",
    "\n",
    "# cg with maximum precision\n",
    "cg_options_prec = Options(tol_abs=1e-20, tol_rel=1e-20, nmax=1e5, results='all')\n",
    "res_cg_prec = cg(A, f, cg_options_prec)\n",
    "print('norm of the residual (cg) = {}'.format(\n",
    "    np.linalg.norm(A@res_cg_prec.x - f)))\n",
    "\n",
    "# gmres with maximum precision\n",
    "gmres_options_prec = Options(tol_abs=1e-20, tol_rel=1e-20, nmax=1e5, results='all')\n",
    "res_gmres_prec = gmres(A, f, gmres_options_prec)\n",
    "print('norm of the residual (gmres) = {}'.format(\n",
    "    np.linalg.norm(A@res_gmres_prec.x - f)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This result is really good. Our algorithms are able to yield results in the order of the ``scipy.sparse.linalg.spsolve`` algorithm and thus convincing precision. Feel free to play around with the relative and absolute tolerances, as well as the maximum number of iterations."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "372px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}