Poisson2D.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Partial Differential Equations</h1>\n",
    "<p>A partial differential equation (PDE) of the function $u(x_1,\\dots,x_n)$ is an equation of the form:\n",
    "\\begin{equation}\n",
    "    F\\big( \\underbrace{x_1,\\dots,x_n}_{\\text{variables}} , \\underbrace{u}_{\\text{function}}, \\underbrace{ \\frac{\\partial u}{\\partial x_1}, \\dots, \\frac{\\partial u}{\\partial x_n}}_{\\text{first derivatives}}, \\underbrace{ \\frac{\\partial^2 u}{\\partial x_1^2}, \\frac{\\partial^2 u}{\\partial x_1\\partial x_2}, \\dots }_{\\text{second derivatives}}, \\underbrace{\\dots}_{\\text{higher order derivatives}} \\big) = 0 \n",
    "\\end{equation}\n",
    "\n",
    "<p> We will focus on linear $2^{nd}$ order PDEs in 2D, which can be written in the form: \n",
    "\\begin{equation}\n",
    "a \\frac{\\partial^2 u}{\\partial x^2} + b \\frac{\\partial^2 u}{\\partial x \\partial y} + c \\frac{\\partial^2 u}{\\partial y^2} + \\underbrace{\\dots}_{\\text{lower terms}} = 0\n",
    "\\end{equation}\n",
    "Like for conic equations, we can define three classes of linear $2^{nd}$ order PDEs, based on the value of discriminant $b^2 - 4 a c $ : \n",
    "    <li> $b^2 - 4 a c < 0$ : <b>elliptic</b>, e.g. the Lapace equation in 2D\n",
    "    \\begin{equation}\n",
    "    \\frac{\\partial^2 u}{\\partial x^2 } + \\frac{\\partial^2 u}{\\partial y^2} = 0\n",
    "    \\end{equation}\n",
    "<li> $b^2 - 4 a c = 0$ : <b>parabolic</b>, e.g. the Diffusion equation in 1D\n",
    "    \\begin{equation}\n",
    "    \\frac{\\partial^2 u}{\\partial x^2} - \\frac{\\partial u}{\\partial t } = 0\n",
    "    \\end{equation}    \n",
    "<li> $b^2 - 4 a c > 0$ : <b>hyperbolic</b>, e.g. the Wave equation in 1D\n",
    "    \\begin{equation}\n",
    "    \\frac{\\partial^2 u}{\\partial x^2} - \\frac{\\partial^2 u}{\\partial t^2} = 0\n",
    "    \\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Poisson equation in 2D</h1>\n",
    "<p>The Poisson equation is a $2^{nd}$ order PDE of the elliptic type: \n",
    "\\begin{equation}\n",
    "\\frac{\\partial^2 u}{\\partial x^2} + \\frac{\\partial^2 u}{\\partial y^2} = f(x,y)\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a 2D electrostatic problem the Poisson equation can be written as\n",
    "\\begin{equation}\n",
    "\\nabla^2 V(x,y) = - \\rho( x, y )\n",
    "\\end{equation}\n",
    "where $\\nabla^2=\\sum_i\\frac{\\partial^2}{\\partial x_i^2}$ is the Laplace operator, $V$ is the potential, and $\\rho$ is the charge density. We note that the Poisson equation differs from the Lapace equation because of the presence of the source term $\\rho$, which makes the equation inhomogeneous. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Assignment</h3>\n",
    "We want to solve the 2D Poisson equation for this source term:\n",
    "\\begin{equation}\n",
    "\\rho(x,y) = sin\\left( 6 \\pi \\frac{x}{L} \\right) sin\\left( 2 \\pi \\frac{y}{L}  \\right)\n",
    "\\end{equation}\n",
    "whith $L=10$. \n",
    "<h3>Summary</h3>\n",
    "We will use three methods: \n",
    "<li> analytical solution\n",
    "<li> spectral method \n",
    "<li> finite difference method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Analytical solution</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p>For the aforementioned source term, it is easy to verify that the solution of the 2D Poisson equation is:\n",
    "\\begin{equation}\n",
    "V(x,y)=\\frac{L^2}{40\\pi^2} sin\\left( 6 \\pi \\frac{x}{L} \\right) sin\\left( 2 \\pi \\frac{y}{L}  \\right)\n",
    "\\end{equation}\n",
    "Let's plot $\\rho(x,y)$ and $V(x,y)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the periodicity\n",
    "L = 10.0\n",
    "# Define the number of points used to discretize the source and the potential on a NxN mesh\n",
    "N = 25\n",
    "\n",
    "# Create the list that contains the x positions \n",
    "x = []\n",
    "for i in range(N) : \n",
    "    x.append(float(i)/float(N)*L)\n",
    "\n",
    "# Create the list that contains the y positions \n",
    "y = []\n",
    "for j in range(N) : \n",
    "    y.append(float(j)/float(N)*L)\n",
    "\n",
    "# Show the arrays\n",
    "print(\"x = \", x)\n",
    "print(\"y = \", y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Access elements of the array x (Python experts can skip)\n",
    "print(\"x[0] = \", x[0]) \n",
    "print(\"x[1] = \", x[1]) \n",
    "print(\"x[N-1] = \", x[N-1]) \n",
    "print(\"x[-1] = \", x[-1]) \n",
    "print(\"x[-N] = \", x[-N]) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create X and Y grids to easily handle gridpoints\n",
    "import numpy as np\n",
    "X, Y = np.meshgrid(x, y)\n",
    "print(\"X = \", X)\n",
    "print(\"Y = \", Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define rho and V on the grid\n",
    "import numpy as np\n",
    "\n",
    "# rho \n",
    "rho = np.empty([N,N]) \n",
    "rho = np.sin(X/L*np.pi*6.0)*np.sin(Y/L*np.pi*2.0)\n",
    "\n",
    "# V analytical solution \n",
    "Van = np.empty([N,N]) \n",
    "Van = (np.sin(X/L*np.pi*6.0)*np.sin(Y/L*np.pi*2.0))/(40.0*np.pi*np.pi/L/L)\n",
    "\n",
    "print(\"rho = \", rho)\n",
    "print(\"Van = \", Van)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# method to plot a 3D function \n",
    "def plot3D( X, Y, Z ) :\n",
    "    from mpl_toolkits.mplot3d import Axes3D\n",
    "    import matplotlib.pyplot as plt\n",
    "    from matplotlib import cm\n",
    "    from matplotlib.ticker import LinearLocator, FormatStrFormatter\n",
    "    import numpy as np\n",
    "\n",
    "\n",
    "    fig = plt.figure()\n",
    "    ax = fig.gca(projection='3d')\n",
    "    \n",
    "    # Plot the surface.\n",
    "    surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,\n",
    "                       linewidth=0, antialiased=False)\n",
    "\n",
    "    # Customize the z axis.\n",
    "    ax.set_zlim(-1.01, 1.01)\n",
    "    ax.zaxis.set_major_locator(LinearLocator(10))\n",
    "    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))\n",
    "\n",
    "    # Set labels\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')\n",
    "    ax.set_zlabel('z')\n",
    "\n",
    "    # Add a color bar which maps values to colors.\n",
    "    fig.colorbar(surf, shrink=0.5, aspect=5)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot3D( X, Y, rho)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot3D( X, Y, Van)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Check for understanding</h3> \n",
    "<i>Q: How many periods are along the x and y direction?</i>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Spectral method</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p>Both the potential and the source are 2D periodic functions, with periodicity $[L,L]$, and can be therefore expanded in plane waves using the following Fourier expansion:\n",
    "\\begin{equation}\n",
    "    \\rho(\\mathbf{r}) = \\sum_\\mathbf{G} \\rho(\\mathbf{G}) e^{i\\mathbf{G}\\cdot\\mathbf{r}}\n",
    "\\end{equation}\n",
    "where the Fourier components $\\rho(\\mathbf{G})$ are obtained as:\n",
    "\\begin{equation}\n",
    "    \\rho(\\mathbf{G}) = \\frac{1}{L^2}\\int d\\mathbf{r} \\rho(\\mathbf{r}) e^{-i\\mathbf{G}\\cdot\\mathbf{r}}\n",
    "\\end{equation}\n",
    "$\\mathbf{r}=(x,y)$, $\\mathbf{G}=\\frac{2\\pi}{L}(n,m)$, with $n=0,\\pm 1, \\pm 2, \\dots $ and $m=0,\\pm 1, \\pm 2, \\dots $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Poisson equation in Fourier space becomes a simple multiplication\n",
    "\\begin{equation}\n",
    "G^2 V(\\mathbf{G}) = \\rho(\\mathbf{G})\n",
    "\\end{equation}\n",
    "where $G^2 = \\mathbf{G}\\cdot\\mathbf{G}=G_x^2+G_y^2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can therefore follow this protocol to obtain $V(x,y)$ : \n",
    "<li> Fourier transform $\\rho(\\mathbf{r}) \\rightarrow \\rho(\\mathbf{G})$\n",
    "<li> Solve the Poisson equation in Fourier space $V(\\mathbf{G})=\\frac{\\rho(\\mathbf{G})}{G^2}$\n",
    "<li> Fourier transform $V(\\mathbf{r}) \\leftarrow V(\\mathbf{G})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define G vectors\n",
    "import numpy as np\n",
    "gx = np.fft.fftfreq(N,1/N) * 2*np.pi/L\n",
    "gy = np.fft.fftfreq(N,1/N) * 2*np.pi/L\n",
    "print(\"list of n\",np.fft.fftfreq(N,1/N))\n",
    "print(\"gx = \",gx)\n",
    "print(\"gy = \",gy)\n",
    "Gx, Gy = np.meshgrid(gx, gy)\n",
    "print(\"Gx = \", Gx)\n",
    "print(\"Gy = \", Gy)\n",
    "G2 = np.empty([N,N])\n",
    "G2 = Gx*Gx + Gy*Gy\n",
    "print(\"G2 = \", G2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fourier transform rho \n",
    "import numpy as np \n",
    "Rrho = rho.astype(complex)\n",
    "Grho = np.fft.fftn(Rrho)\n",
    "print(Grho)\n",
    "print(\"Grho.real = \",Grho.real)\n",
    "print(\"Grho.imag = \",Grho.imag)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Check for understanding</h3> \n",
    "<i>Q: Why do we need to work with complex numbers?</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analize Grho\n",
    "for j in range(N) : \n",
    "    for i in range(N) : \n",
    "        if (abs(Grho[i,j].real) > 0.00001) : \n",
    "            print('real:',Grho[i,j].real,'n:',Gx[i,j]*L/(2.0*np.pi),'m:',Gy[i,j]*L/(2.0*np.pi))\n",
    "        if (abs(Grho[i,j].imag) > 0.00001) : \n",
    "            print('imag:',Grho[i,j].imag,'n:',Gx[i,j]*L/(2.0*np.pi),'m:',Gy[i,j]*L/(2.0*np.pi))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Check for understanding</h3> \n",
    "<i>Q: Can you explain why we have both positive and negative values for $n$ and $m$?</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve Poisson equation in Fourier space\n",
    "GV = np.empty([N,N],dtype=complex)\n",
    "GV[0,0] = 0.0 # arbitrary constant\n",
    "for j in range(N) : \n",
    "    for i in range(N) :\n",
    "        if( i != 0 or j!=0 ) : \n",
    "            GV[i,j] = Grho[i,j] / G2[i,j]\n",
    "print(\"GV = \",GV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Check for understanding</h3> \n",
    "<i>Q: Why is $V(\\mathbf{G}=0)$ arbitrary?</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fourier transform V to real space\n",
    "RV = np.fft.ifftn(GV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot3D(X,Y,RV.real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the max Abs difference between A[i,j] and B[i,j]\n",
    "def maxAbsDiff(A,B) : \n",
    "    import numpy as np\n",
    "    C = A-B\n",
    "    Cflat = np.abs(C.flatten())\n",
    "    return max(Cflat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxAD = maxAbsDiff(RV.real,Van)\n",
    "print(\"maxAD = \",maxAD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Finite difference</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second derivative of a function $f$ can be computed numerically using the central difference formula\n",
    "\\begin{equation}\n",
    "\\frac{d^2f}{dx^2} \\approx \\frac{f(x+h)+f(x-h)-2f(x)}{h^2}\n",
    "\\end{equation}\n",
    "where $h$ is a small increment. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a similar way, we can discretize the Lapace operator on the 2D grid:\n",
    "\\begin{equation}\n",
    "\\nabla^2 V(x,y) \\approx \\frac{V(x+\\Delta,y)+V(x-\\Delta,y)+V(x,y+\\Delta)+V(x,y-\\Delta)-4V(x,y)}{\\Delta^2} \n",
    "\\end{equation}\n",
    "This yields the discretized form of the Poisson equation: \n",
    "\\begin{equation}\n",
    "\\nabla^2V_{i,j} = \\frac{V_{i+1,j}+V_{i-1,j}+V_{i,j+1}+V_{i,j-1}-4V_{i,j}}{\\Delta^2} =-\\rho_{i,j}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can interpret this as a self-consistent equation: the value of $V$ in every point of the grid depends on the average of its four neighbors plus the source contribution.\n",
    "\\begin{equation}\n",
    "V^{n+1}_{i,j} = \\frac{V^n_{i+1,j}+V^n_{i-1,j}+V^n_{i,j+1}+V^n_{i,j-1}}{4} + \\frac{\\Delta^2}{4} \\rho_{i,j}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We solve the Possion equation by iterative finite differences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Delta\n",
    "Delta = x[1]-x[0]\n",
    "print(\"Delta = \", Delta )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Method that plots values in listA and optionally of listB\n",
    "def plotMaxDiff(listA,listB=[]) : \n",
    "    import matplotlib.pyplot as plt\n",
    "    xB = []\n",
    "    for i in range(len(listB)) : \n",
    "        xB.append(i)\n",
    "    plt.plot(xB,listB,'bo-')\n",
    "    xA = []\n",
    "    for i in range(len(listA)) : \n",
    "        xA.append(i)\n",
    "    plt.plot(xA,listA,'ro-')\n",
    "    plt.xlabel('iteration')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define how Vnew is obtained from Vold \n",
    "def updateV(Vold,source,N,Delta) : \n",
    "    import numpy as np\n",
    "    Vnew = np.empty([N,N]) \n",
    "    for j in range(N) : \n",
    "        for i in range(N) : \n",
    "            iplus = i+1\n",
    "            if (iplus==N) : \n",
    "                iplus=0\n",
    "            jplus = j+1\n",
    "            if (jplus==N) : \n",
    "                jplus=0\n",
    "            Vnew[i,j]=0.25*(Vold[i-1,j]+Vold[iplus,j]+Vold[i,j-1]+Vold[i,jplus])+ 0.25*Delta*Delta*source[i,j]\n",
    "    return Vnew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init\n",
    "list_max = []\n",
    "import numpy as np\n",
    "Vold = np.zeros([N,N])\n",
    "Vnew = Vold\n",
    "plot3D(X,Y,Vnew)\n",
    "list_max.append( maxAbsDiff(Vnew,Van) )\n",
    "plotMaxDiff(list_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#iterate\n",
    "Vold = Vnew \n",
    "Vnew = updateV(Vold,rho,N,Delta)\n",
    "plot3D(X,Y,Vnew)\n",
    "list_max.append( maxAbsDiff(Vnew,Van) )\n",
    "plotMaxDiff(list_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We introduce a small variation to the previous method. Here self-consistency is obtained by mixing at every iteration the old and new values: \n",
    "\\begin{equation}\n",
    "V^{n+1}_{i,j} = (1-\\omega)V^{n}_{i,j} +\\omega \\left[\\frac{V^n_{i+1,j}+V^n_{i-1,j}+V^n_{i,j+1}+V^n_{i,j-1}}{4} + \\frac{\\Delta^2}{4} \\rho_{i,j}\\right]\n",
    "\\end{equation}\n",
    "where $\\omega$ is a parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def newUpdateV(Vold,source,N,Delta,omega) : \n",
    "    import numpy as np\n",
    "    Vnew = np.empty([N,N]) \n",
    "    for j in range(N) : \n",
    "        for i in range(N) : \n",
    "            iplus = i+1\n",
    "            if (iplus==N) : \n",
    "                iplus=0\n",
    "            jplus = j+1\n",
    "            if (jplus==N) : \n",
    "                jplus=0\n",
    "            Vnew[i,j]=(1.0-omega)*Vold[i,j]+omega*0.25*(Vold[i-1,j]+Vold[iplus,j]+Vold[i,j-1]+Vold[i,jplus])+omega*0.25*Delta*Delta*source[i,j]\n",
    "    return Vnew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#init \n",
    "list_max_new = []\n",
    "import numpy as np\n",
    "Vold = np.zeros([N,N])\n",
    "Vnew = Vold\n",
    "plot3D(X,Y,Vnew)\n",
    "list_max_new.append( maxAbsDiff(Vnew,Van) )\n",
    "plotMaxDiff(list_max_new,list_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# iterate\n",
    "Vold = Vnew \n",
    "Vnew = newUpdateV(Vold,rho,N,Delta,2)\n",
    "plot3D(X,Y,Vnew)\n",
    "list_max_new.append( maxAbsDiff(Vnew,Van) )\n",
    "plotMaxDiff(list_max_new,list_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Check for understanding</h3> \n",
    "<i>Q: Why does the new method improve convergence?</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}