{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TD8: Matrix inversion optimization around an example\n",
    "\n",
    "Solving linear systems $AX = B$ is generally a crucial step when solving PDEs. The choice of the right numerical method to do that can have a major impact on the performance of an algorithm both in terms of computational cost and memory cost. The structure and the properties of $A$ should determine the method.\n",
    "\n",
    "In this TD, we will adress this issue through the example of a familiar problem: the 2D Laplace problem on a squared domain. The goal is to solve in an efficient way:\n",
    "$$\n",
    "    \\begin{cases}\n",
    "        \\Delta f = 0 \\text{ on } \\Omega\\\\\n",
    "        f = g \\text{ on } \\partial\\Omega\n",
    "    \\end{cases}\n",
    "$$\n",
    "with  $\\Omega = ]-\\frac{L}{2}, \\frac{L}{2}[^2$. In TD4, we used the 5-point stencil $\\Delta_5$ to approximate the laplacian operator with finite-differences. We recall the definition of this stencil:\n",
    "$$\n",
    "\\Delta_5 f (x, y) = \\frac1{h^2}\\left(f \\left(x + h, y\\right) + f \\left(x - h, y\\right) + f \\left(x, y +h\\right) + f\\left(x, y -h\\right) - 4f\\left(x, y\\right)\\right)\n",
    "$$\n",
    "\n",
    "We also recall that we typically solve this problem in three steps:\n",
    "1. Define a grid which covers $\\Omega\\cup\\partial\\Omega$ and is suitable for the chosen solving method.\n",
    "2. Build the matrix which constrains the solution on the domain.\n",
    "3. Find the inverse of this matrix to solve the problem.\n",
    "\n",
    "We will first need the following Python packages for this TD:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To draw matplotlib plots within this notebook.\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.linalg as la"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: A naive algorithm\n",
    "\n",
    "In TD4, we used a spatial grid $\\mathcal{G}$ of size $N\\times N$ defined by:\n",
    "$$\n",
    "    \\mathcal{G} = \\{(x_i = -L/2 + ih, y_j = -L/2 + jh), i = 0, ..., N - 1 \\text{ and } j = 0, ..., N-1\\} \n",
    "$$\n",
    "with $h = \\frac{L}{N-1}$ the spatial step.\n",
    "\n",
    "On this grid, the 2D Laplace problem was rephrased in terms of two distinct set of constraints:\n",
    "1. $\\forall (x_i, y_j) \\in \\mathcal{G} \\cap \\Omega, \\Delta_5 f(x_i, y_j) = 0$\n",
    "2. $\\forall (x_i, y_j) \\in \\mathcal{G} \\backslash \\Omega, f(x_i, y_j) = g(x_i, y_j)$\n",
    "\n",
    "We turned these set of constraints into a linear system $AX=B$ where $A$ is a matrix of shape $(N \\times N, N\\times N)$ and $B$ is a column matrix of shape $(N \\times N, 1)$, and we finally solved it using a LU solver namely  [scipy.linalg.lu_factor](https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.lu_factor.html) and [scipy.linalg.lu_solve](https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.lu_solve.html) functions.\n",
    "\n",
    "The following code implements this algorithm with $g(x,y)=\\cos (2\\pi(2x + y)/L)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matrix_builder(N, L):\n",
    "    x = np.linspace(-L/2, L/2, N)\n",
    "    y = np.linspace(-L/2, L/2, N)\n",
    "    \n",
    "    # We define the grid with meshgrid function with a matrix indexing ('ij' parameter) so that distinct matrix rows correspond to distinct x coordinates\n",
    "    # Similarly, distinct matrix columns also correspond to distinct y coordinates. \n",
    "    x_grid, y_grid = np.meshgrid(x, y, indexing='ij')\n",
    "    \n",
    "    A = np.zeros((N * N, N * N))\n",
    "    B = np.zeros(N * N)\n",
    "    \n",
    "    def g(x, y):\n",
    "        return np.cos(2*np.pi/L * (2*x + y))\n",
    "    \n",
    "    for i in range(x_grid.shape[0]):\n",
    "        for j in range(x_grid.shape[1]):\n",
    "            if np.abs(x_grid[i, j]) < L/2 and np.abs(y_grid[i, j]) < L/2: # We are inside \\Omega\n",
    "                A[i*N + j, i*N + j] = -4.0\n",
    "                A[i*N + j, (i - 1)*N + j] = 1.0\n",
    "                A[i*N + j, (i + 1)*N + j] = 1.0\n",
    "                A[i*N + j, i*N + j + 1] = 1.0\n",
    "                A[i*N + j, i*N + j - 1] = 1.0\n",
    "            else: # We are outside \\Omega\n",
    "                A[i*N + j, i*N + j] = 1.0\n",
    "                B[i*N + j] = g(x_grid[i, j], y_grid[i, j])\n",
    "                \n",
    "    return A, B\n",
    "\n",
    "def matrix_solver(A, B):\n",
    "    return la.lu_solve(la.lu_factor(A), B)\n",
    "\n",
    "def plot_solution(N, L, X):\n",
    "    h = L/(N - 1) # Becareful, h is not equal to L / N\n",
    "\n",
    "    # We reshape X vector into a (N, N) matrix\n",
    "    # Row indices correspond to different x coordinates, and columns indices to different y coordinates\n",
    "    z_grid = X.reshape((N, N))\n",
    "    \n",
    "    # We will pay attention to the documentation of pcolor function to plot an accurate solution\n",
    "    # We need: to avoid mixing x and y axis\n",
    "    # We need: to define a grid of corner points instead of using directly x_grid and y_grid variables\n",
    "    plt.figure()\n",
    "    x_corner_points = np.linspace(-L/2 - h/2, L/2 + h/2, N + 1)\n",
    "    y_corner_points = np.linspace(-L/2 - h/2, L/2 + h/2, N + 1)\n",
    "    plt.pcolor(x_corner_points, y_corner_points, z_grid.T, cmap='jet')\n",
    "    plt.colorbar()\n",
    "    plt.xlabel('$x$')\n",
    "    plt.ylabel('$y$')\n",
    "    plt.title(\"Laplace equation with a 5-point stencil\")\n",
    "    plt.show()\n",
    "    \n",
    "def naive_solver(N, L, plot=False):\n",
    "    A, B = matrix_builder(N, L)\n",
    "    X = matrix_solver(A, B)\n",
    "    if plot:\n",
    "        plot_solution(N, L, X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**1)** Using the previous *naive_solver* function, solve this 2D Laplace problem for arbitrary values of $N$ and $L$. We will also plot the solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2)** What is the most computationally expensive part of the previous code in terms of CPU operations? Deduce from this the complexity of this algorithm as a function of $N$.\n",
    "\n",
    "**3)** How should evolve the memory cost of this algorithm with $N$?\n",
    "\n",
    "**4)** Measure the runtime of the previous algorithm for a reasonable range of $N$ values. We will plot our measures on a log-log scale. We will use successive calls of [time.time](https://docs.python.org/3/library/time.html#time.time) function to measure the runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# To measure the runtime of your code, typically we do:\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "# my algorithm ....\n",
    "\n",
    "elapsed_time = time.time() - start_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**5)** You should observe an asymptotic linear trend on the previous plot. Using [scipy.stats.linregress](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.linregress.html) measure the slope of this linear trend. Is this what you were expecting?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.stats as stats\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**6)** We have been solving the linear system $AX=B$ with a common LU factorization, but there are much smarter ways to do that in this specific case. What can you tell about the structure of the $A$ matrix we built with *matrix_builder* function? We will visualize $A$ matrix with [pyplot.matshow](https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.matshow.html) function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Sparse representations of matrices\n",
    "\n",
    "The previous matrix $A$ is a banded matrix with 4 non-zero diagonals only, meaning that the number of non zero coefficients is highly negligible compared to the number of zero coefficients when $N$ is big enough. We call such a matrix a *sparse matrix* and we certainly can make use of that outstanding property...\n",
    "\n",
    "The [scipy.sparse](https://docs.scipy.org/doc/scipy/reference/sparse.html) package gives us handy functions to define and manipulate these sparse matrices.\n",
    "\n",
    "**7)** Build a tridiagonal matrix $A$ with coefficients on the main diagonal equal to 1, and coefficients on the first diagonals above and below the main diagonal equal to 2. We will use [scipy.sparse.diags](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.diags.html#scipy.sparse.diags) function to build this matrix, and we will plot it with [pyplot.matshow](https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.matshow.html) function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse as sparse\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**8)** Now let us go back to our 2D Laplace problem. On the model of the code of part 1, define the following new set of functions:\n",
    "* *sparse_matrix_builder* function which builds $A$ and $B$ matrices using sparse data structures from *scipy.sparse* package (at least for $A$ matrix)\n",
    "* *sparse_matrix_solver* function which solves the linear system $AX=B$ for sparse matrices using [scipy.sparse.linalg.spsolve](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.spsolve.html#scipy.sparse.linalg.spsolve) function\n",
    "* *sparse_solver* function which actually solves our 2D Laplace problem using these previous functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse.linalg as splinalg\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**9)** Now, how should evolve the memory cost of this algorithm with $N$?\n",
    "\n",
    "**10)** Measure the time complexity of this new algorithm on the model of what we did in the first part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
