PSPython_02-pseudospectral-collocation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import HTML\n",
    "css_file = './custom.css'\n",
    "HTML(open(css_file, \"r\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Content provided under a Creative Commons Attribution license, CC-BY 4.0; code under MIT License. (c)2015 [David I. Ketcheson](http://davidketcheson.info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Version 0.3 - May 2022"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "font = {'size'   : 15}\n",
    "matplotlib.rc('font', **font)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pseudospectral collocation methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to lesson 2 -- in which we'll learn about the main focus of this course: pseudospectral collocation methods for wave equations.  But first, a bit of review."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advection-diffusion (again)\n",
    "\n",
    "$$\n",
    "\\newcommand{\\F}{\\mathcal F}\n",
    "\\newcommand{\\Finv}{{\\mathcal F}^{-1}}\n",
    "$$\n",
    "\n",
    "In notebook 1 we solved the scalar, linear 1D evolution equation:\n",
    "\n",
    "$$u_t + u_x = \\epsilon u_{xx}$$\n",
    "\n",
    "by taking the Fourier transform, which \"diagonalizes\" this infinite-dimensional dynamical system:\n",
    "\n",
    "$$\\hat{u}_t + i\\xi\\hat{u} = -\\xi^2 \\epsilon \\hat{u}$$\n",
    "\n",
    "with solution\n",
    "\n",
    "$$\\hat{u}(t) = e^{(-i \\xi - \\epsilon \\xi^2)t} \\hat{u}(0)$$\n",
    "\n",
    "for each mode.\n",
    "\n",
    "To construct the full solution, we simply take the inverse Fourier transform.  All together, this algorithm looks like:\n",
    "\n",
    "$$\n",
    "u(t) = \\Finv \\left(D\\left[e^{(-i \\xi - \\epsilon \\xi^2)t}\\right]\\F(u) \\right),\n",
    "$$\n",
    "\n",
    "where we have written $D[f(\\xi)]$ to denote the diagonal matrix whose $(j,j)$ entry is given by $f(\\xi_j)$.\n",
    "In the exact solution, the wavenumbers $\\xi$ range from $-\\infty$ to $+\\infty$ (and $D$ is infinite), but in practice we compute on a finite interval of length $L$ with $m$ collocation points.\n",
    "The wavenumbers are then given by the formula\n",
    "\n",
    "\\begin{align}\n",
    "\\xi_j & = \\frac{2 \\pi j}{L} & \\text{for } -m/2 \\le j < m/2,\n",
    "\\end{align}\n",
    "\n",
    "although the FFT routine orders them differently."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make the algorithm even more explicit by recognizing that the FFT (and its inverse) is a linear map, so it can be represented as a matrix multiplication.  Thus\n",
    "\n",
    "$$\\F(u) = Fu$$\n",
    "\n",
    "where $F$ is a certain $m \\times m$ matrix.  This matrix has a number of interesting properties (see e.g. Trefethen, Ch. 3), but for the moment we are only interested in the fact that it is a linear operator.  We can reverse engineer it by applying $\\F$ to the columns of the identity matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def F_matrix(m):\n",
    "    F = np.zeros((m,m),dtype=complex)\n",
    "    for j in range(m):\n",
    "        v = np.zeros(m)\n",
    "        v[j] = 1.\n",
    "        F[:,j] = np.fft.fft(v)\n",
    "    return F\n",
    "\n",
    "print(F_matrix(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that $F$ is symmetric; this is true for any value of $m$.  Also, $F$ is essentially unitary -- that is, it is possible to choose a normalization such that $F$ is unitary, but the normalization in common use means that $FF^*=mI$ (and $F^{-1}(F^{-1})^* = m^{-1} I$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We won't work with this matrix in practice since applying it naively is much slower than the FFT.  But it is a useful representation.  Thus we can write our \"algorithm\" simply as\n",
    "\n",
    "$$\n",
    "u(t) = F^{-1} D\\left[e^{(-i \\xi - \\epsilon \\xi^2)(t-t_0)}\\right] F u(t_0).\n",
    "$$\n",
    "\n",
    "Now it is completely clear that we are simply applying a similarity transformation that diagonalizes the dynamics of our system:\n",
    "\n",
    "$$\n",
    "F u(t) =  D\\left[e^{(-i \\xi - \\epsilon \\xi^2)(t-t_0)}\\right] F u(t_0).\n",
    "$$\n",
    "\n",
    "We can solve any linear, scalar constant-coefficient 1D evolution equation (with periodic boundary conditions) in similar manner, by simply replacing the argument of $D$.  Note that this algorithm is:\n",
    "\n",
    "- Exact in time (the only error is due to our initial truncation of the Fourier series by sampling the data at a finite set of points); and\n",
    "- Unconditionally stable.\n",
    "\n",
    "Here is the implementation as a reminder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spectral_representation(x0,uhat):\n",
    "    u_fun = lambda y : np.real(np.sum(uhat*np.exp(1j*xi*(y+x0))))/len(uhat)\n",
    "    u_fun = np.vectorize(u_fun)\n",
    "    return u_fun\n",
    "\n",
    "# Spatial grid\n",
    "m=64                            # Number of grid points in space\n",
    "L = 2 * np.pi                   # Width of spatial domain\n",
    "x = np.arange(-m/2,m/2)*(L/m)   # Grid points\n",
    "dx = x[1]-x[0]                  # Grid spacing\n",
    "\n",
    "# Temporal grid\n",
    "tmax=4.0     # Final time\n",
    "N = 25       # number grid points in time\n",
    "k = tmax/N   # interval between output times\n",
    "\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L  # Wavenumber \"grid\"\n",
    "# (this is the order in which numpy's FFT gives the frequencies)\n",
    "\n",
    "# Initial data\n",
    "u = np.sin(2*x)**2 * (x<-L/4)\n",
    "uhat0 = np.fft.fft(u)\n",
    "\n",
    "epsilon=0.01  # Diffusion coefficient\n",
    "a = 1.0       # Advection coefficient\n",
    "\n",
    "# Store solutions in a list for plotting later\n",
    "frames = [u.copy()]\n",
    "ftframes = [uhat0.copy()]\n",
    "\n",
    "# Now we solve the problem\n",
    "for n in range(1,N+1):\n",
    "    t = n*k\n",
    "    uhat = np.exp(-(1.j*xi*a + epsilon*xi**2)*t) * uhat0\n",
    "    u = np.real(np.fft.ifft(uhat))\n",
    "    frames.append(u.copy())\n",
    "    ftframes.append(uhat.copy())\n",
    "    \n",
    "# Set up plotting\n",
    "fig = plt.figure(figsize=(9,4)); axes = fig.add_subplot(111)\n",
    "line, = axes.plot([],[],lw=3)\n",
    "axes.set_xlim((x[0],x[-1])); axes.set_ylim((0.,1.))\n",
    "plt.close()\n",
    "\n",
    "x_fine = np.linspace(x[0],x[-1],1000)\n",
    "\n",
    "def plot_frame(i):\n",
    "    uhat = ftframes[i]\n",
    "    u_spectral = spectral_representation(x[0],uhat)\n",
    "    line.set_data(x_fine,u_spectral(x_fine));\n",
    "    #line.set_data(x,frames[i])\n",
    "    axes.set_title('t='+str(i*k))\n",
    "\n",
    "# Animate the solution\n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                   frames=len(frames),\n",
    "                                   interval=200)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable coefficients\n",
    "\n",
    "Things become more interesting as soon as we introduce spatially-varying coefficients or nonlinearity.  In either case, the Fourier transform no longer provides a global diagonalization -- instead, it can be thought of as diagonalizing the system at one instant in time.  That means that we will have to discretize in time, and the time step we can use will be restricted by accuracy and (possibly) stability considerations.\n",
    "\n",
    "Consider now the **variable-coefficient advection equation**\n",
    "\n",
    "$$ u_t + a(x) u_x = 0.$$\n",
    "\n",
    "In a purely spectral method, we would take the Fourier transform of this equation and we would need to compute a convolution of $a(x)$ with $u$, which is computationally quite expensive relative to the rest of what we will do.  Specifically, the convolution would require ${\\mathcal O}(m^2)$ operations, whereas the FFT requires only ${\\mathcal O}(m \\log m)$.  To avoid this expense, we use the \"pseudospectral\" approach: we compute the derivative $u_x$ in the spectral way:\n",
    "\n",
    "$$ u_x = F^{-1} \\left(D\\left[i\\xi\\right] F u\\right),$$\n",
    "\n",
    "and then multiply by $a(x)$ in physical space:\n",
    "\n",
    "$$ (a(x) u_x)_j = a(x_j) \\left(F^{-1}  \\left(D\\left[i\\xi\\right] F u\\right) \\right)_j.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Letting\n",
    "$$\n",
    "    A = \\begin{pmatrix} a(x_1) \\\\ & a(x_2) \\\\ & & \\ddots \\\\ & & & a(x_m) \\end{pmatrix},\n",
    "$$\n",
    "\n",
    "we have the semi-discrete system\n",
    "\n",
    "$$\n",
    "U'(t) = -A \\Finv \\left(D\\left[i\\xi\\right] \\F(U)\\right).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time discretization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Next we need to integrate in time.  To choose an appropriate time integrator, we would like to know the spectrum of our semi-discretization.  We can compute it explicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have a linear ODE system of the form $U'(t) = MU(t)$, where\n",
    "\n",
    "$$M = -AF^{-1}D[i\\xi]F.$$\n",
    "\n",
    "The code below computes the spectrum of $M$ for a few choices of $a(x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 32\n",
    "F = F_matrix(m)\n",
    "Finv = np.linalg.inv(F)\n",
    "L = 2 * np.pi\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "D = np.diag(1.j*xi)\n",
    "x = np.arange(-m/2,m/2)*(L/m)\n",
    "\n",
    "M = np.dot(Finv,np.dot(D,F))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_spectrum(a,m=64):\n",
    "    F = F_matrix(m)\n",
    "    Finv = np.linalg.inv(F)\n",
    "    L = 2 * np.pi\n",
    "    xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "    D = np.diag(1.j*xi)\n",
    "    x = np.arange(-m/2,m/2)*(L/m)\n",
    "    A = np.diag(a(x))\n",
    "    M = -np.dot(A,np.dot(Finv,np.dot(D,F)))\n",
    "    lamda = np.linalg.eigvals(M)\n",
    "    print(np.max(np.abs(lamda)))\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8,4),\n",
    "                                   gridspec_kw={'width_ratios': [3, 1]})\n",
    "    ax1.plot(x,a(x)); ax1.set_xlim(x[0],x[-1])\n",
    "    ax1.set_xlabel('$x$'); ax1.set_ylabel('$a(x)$')\n",
    "    ax2.plot(np.real(lamda),np.imag(lamda),'ko')\n",
    "    ax2.axis('equal')\n",
    "    ax2.set_xlabel('$Re(\\lambda)$')\n",
    "    ax2.set_ylabel('$Im(\\lambda)$')\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = lambda x : np.ones(len(x))\n",
    "plot_spectrum(a,m=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = lambda x : 2 + np.sin(x)\n",
    "plot_spectrum(a,m=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the eigenvalues of $M$ are purely imaginary, regardless of the choice of $a(x)$.  This is not surprising, since $M$ is the product of a real diagonal matrix $A$ and a matrix that is similar to $D$.\n",
    "\n",
    "It is also straightforward to see that the largest eigenvalues of $M$ have magnitude equal to $\\max_i |a(x_i)| m/2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus we should choose a time integrator whose absolute stability region includes part of the imaginary axis -- ideally, a large part of it.  A simple integrator of this type is the explicit midpoint method (also referred to as the *leapfrog* method).  For our problem it amounts to\n",
    "\n",
    "$$\n",
    "U^{n+1} = U^{n-1} + 2\\Delta t M U^n.\n",
    "$$\n",
    "\n",
    "Since this is a 2-step method, we need some other way to take the first step.  For that, we use the explicit Euler method:\n",
    "\n",
    "$$\n",
    "U^{n+1} = U^n + \\Delta t M U^n.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#a = lambda x : np.ones(len(x))\n",
    "a = lambda x : 2 + np.sin(x)\n",
    "\n",
    "m = 64\n",
    "F = F_matrix(m)\n",
    "Finv = np.linalg.inv(F)\n",
    "L = 2 * np.pi\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "D = np.diag(1.j*xi)\n",
    "x = np.arange(-m/2,m/2)*(L/m)\n",
    "A = np.diag(a(x))\n",
    "#M = -np.dot(A,np.dot(Finv,np.dot(D,F)))\n",
    "M = -A@Finv@D@F\n",
    "\n",
    "# Initial data\n",
    "u = np.sin(2*x)**2 * (x<-L/4)\n",
    "dx = x[1]-x[0]\n",
    "dt = 2.0/m/np.max(np.abs(a(x)))\n",
    "#dt = 1./86.73416328005729 + 1e-4\n",
    "T = 10.\n",
    "N = int(np.round(T/dt))\n",
    "\n",
    "frames = [u.copy()]\n",
    "skip = N//100\n",
    "\n",
    "# Start with an explicit Euler step\n",
    "u_new = u + dt*np.dot(M,u)\n",
    "\n",
    "# Now we solve the problem\n",
    "for n in range(1,N+1):\n",
    "    t = n*dt\n",
    "    u_old = u.copy()\n",
    "    u = u_new.copy()\n",
    "    u_new = np.real(u_old + 2*dt*np.dot(M,u))\n",
    "    if ((n % skip) == 0):\n",
    "        frames.append(u_new.copy())\n",
    "    \n",
    "# Set up plotting\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8,8), sharex=True)\n",
    "\n",
    "line1, = ax1.plot([],[],lw=3)\n",
    "ax1.set_xlim((x[0],x[-1])); ax1.set_ylim((-0.1,1.1))\n",
    "ax2.plot(x,a(x),lw=3); ax2.set_ylim(0,3.1)\n",
    "plt.close()\n",
    "\n",
    "def plot_frame(i):\n",
    "    line1.set_data(x,frames[i])\n",
    "    ax1.set_title('t='+str(i*skip*dt))\n",
    "\n",
    "# Animate the solution\n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                   frames=len(frames),\n",
    "                                   interval=200)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: absolute stability\n",
    "The region of absolute stability for the midpoint method is the interval $S=[-i,i]$.  The numerical solution will be stable only if $\\lambda \\Delta t \\in S$ for all eigenvalues $\\lambda$ of $M$.  \n",
    "\n",
    "1. Determine the maximum stable step size.  What happens if you take a step size larger than this value?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, in addition to truncating the Fourier series of the initial condition, we are truncating the Taylor series of the evolution in time.  This leads to truncation error that accumulates, leading to a solution that is less accurate as time progresses.  \n",
    "\n",
    "2. Set $\\Delta t$ equal to the largest stable step size and notice the oscillations that appear in front of the pulse as it propagates.  Try reducing the size of $\\Delta t$ to decrease the amount of truncation error.  Try some different velocity functions $a(x)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A nonlinear wave equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let us consider the case in which the speed of propagation is given by the solution value $u$ itself:\n",
    "\n",
    "$$\n",
    "    u_t + u u_x = 0.\n",
    "$$\n",
    "\n",
    "We will again use the pseudospectral approach to approximate the term $uu_x$, by differentiating in frequency space and then multiplying in physical space.\n",
    "\n",
    "We can implement this easily, using the code above but replacing the entries in the matrix $A$ by the values of $u$ on the grid.  We need to construct this matrix at each time step now, since $u$ is changing in time.\n",
    "\n",
    "This implementation is very inefficient but since we are working on a small grid in one dimension we can afford to ignore efficiency for now.  Run the code below.  Can you explain what happens?  Can you \"fix\" the problem that occurs by adjusting the parameters in the code?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 64\n",
    "F = F_matrix(m)\n",
    "Finv = np.linalg.inv(F)\n",
    "L = 2 * np.pi\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "D = np.diag(1.j*xi)\n",
    "x = np.arange(-m/2,m/2)*(L/m)\n",
    "\n",
    "\n",
    "# Initial data\n",
    "u = np.sin(2*x)**2 * (x<-L/4) + 1.\n",
    "dx = x[1]-x[0]\n",
    "dt = 2.0/m/np.max(np.abs(a(x)))/2.\n",
    "T = 5.\n",
    "N = int(np.round(T/dt))\n",
    "\n",
    "ftframes = [np.fft.fft(u)]\n",
    "skip = N//100\n",
    "\n",
    "# Start with an explicit Euler step\n",
    "A = np.diag(u)\n",
    "M = -np.dot(A,np.dot(Finv,np.dot(D,F)))\n",
    "u_new = u + dt*np.dot(M,u)\n",
    "\n",
    "# Now we solve the problem\n",
    "for n in range(1,N+1):\n",
    "    t = n*dt\n",
    "    u_old = u.copy()\n",
    "    u = u_new.copy()\n",
    "    \n",
    "    A = np.diag(u)\n",
    "    M = -np.dot(A,np.dot(Finv,np.dot(D,F)))\n",
    "    u_new = np.real(u_old + 2*dt*np.dot(M,u))\n",
    "    if ((n % skip) == 0):\n",
    "        ftframes.append(np.fft.fft(u_new))\n",
    "    \n",
    "# Set up plotting\n",
    "fig, ax1 = plt.subplots(1, 1, figsize=(8,4))\n",
    "\n",
    "line1, = ax1.plot([],[],lw=3)\n",
    "ax1.set_xlim((x[0],x[-1])); ax1.set_ylim((0.,2.1))\n",
    "plt.close()\n",
    "\n",
    "def plot_frame(i):\n",
    "    uhat = ftframes[i]\n",
    "    u_spectral = spectral_representation(x[0],uhat)\n",
    "    line1.set_data(x_fine,u_spectral(x_fine));\n",
    "    ax1.set_title('t='+str(i*skip*dt))\n",
    "\n",
    "# Animate the solution\n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                   frames=len(ftframes),\n",
    "                                   interval=200)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  },
  "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": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}