PSPython_03-FFT-aliasing-filtering.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": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation\n",
    "from IPython.display import HTML\n",
    "font = {'size'   : 15}\n",
    "matplotlib.rc('font', **font)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The FFT, aliasing, and filtering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to lesson 3.  Here we'll learn about the Fast Fourier Transform, which we've been using all along.  We'll also learn about a numerical pathology of pseudospectral methods (known as *aliasing*) and one way to avoid it (known as *filtering* or *dealiasing*)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The fast Fourier transform\n",
    "\n",
    "We won't go into great detail regarding the FFT algorithm, since there is already an excellent explanation of the Fast Fourier Transform in Jupyter Notebook form available on the web:\n",
    "\n",
    "- [Understanding the FFT Algorithm](https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/) by Jake Vanderplas\n",
    "\n",
    "Suffice it to say that the FFT is a fast algorithm for computing the discrete Fourier transform (DFT):\n",
    "\n",
    "$$\n",
    "\\hat{u}_\\xi = \\sum_{j=0}^{m-1} u_j e^{-2\\pi i \\xi j/m}\n",
    "$$\n",
    "\n",
    "or its inverse.  The DFT, as we know, is linear and can be computed by multiplying $u$ by a certain $m\\times m$  dense matrix $F$.  Multiplication by a dense matrix requires ${\\mathcal O}(m^2)$ operations.\n",
    "The FFT is a shortcut to compute that matrix-vector product in just ${\\mathcal O}(m \\log m)$ operations by taking advantage of the special structure of $F$.  \n",
    "\n",
    "This is very important for pseudospectral methods, since most of the computational work occurs in computing the Fourier transform and its inverse.  It's also important that we make use of a compiled version of the FFT, since native Python code is relatively slow.  The `np.fft` function is an interface to a compiled FFT library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ordering of wavenumbers in FFT output\n",
    "The vector returned by `np.fft` contains the Fourier coefficients of $u$ corresponding to the wavenumbers\n",
    "\n",
    "$$\n",
    "\\frac{2\\pi}{L} \\{-m/2, -m/2 + 1, \\dots, m/2 - 1\\}.\n",
    "$$\n",
    "\n",
    "However, for computational efficiency the output vector does not use the natural ordering above.  The ordering it uses can be obtained with the following command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 16\n",
    "L = 2*np.pi\n",
    "\n",
    "xi=np.fft.fftfreq(m)*m/(L/(2*np.pi))\n",
    "print(xi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the return vector starts with the nonnegative wavenumbers, followed by the negative wavenumbers.  It may seem strange to you that the range of wavenumbers returned is not symmetric; in the case above, it includes $-8$ but not $+8$.  This apparent asymmetry can be explained once one understands the phenomenon known as *aliasing*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aliasing\n",
    "\n",
    "A numerical grid has a limited resolution.  If you try to represent a rapidly-oscillating function with relatively few grid points, you will observe an effect known as **aliasing**.  This naturally limits the range of frequencies that can be modelled on a given grid.  It can also lead to instabilities in pseudospectral simulations, when generation of high frequencies leads to buildup of lower-frequency energy due to aliasing.\n",
    "\n",
    "The code below plots a sine wave of a given wavenumber $\\xi$, along with its representation on a grid with $m$ points.  Try changing $\\xi$ and notice how for $  |\\xi| \\ge m$ the function looks like a lower-frequency mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import widgets\n",
    "from ipywidgets import interact, interactive\n",
    "\n",
    "def plot_sine(wavenumber=4,grid_points=12,plot_sine=True):\n",
    "    \"Plot sin(2*pi*p), sampled at m equispaced points.\"\n",
    "    x  = np.linspace(0,1,grid_points+1); # grid\n",
    "    xf = np.linspace(0,1,1000) # fine grid\n",
    "    y  = np.sin(wavenumber*np.pi*x)\n",
    "    yf = np.sin(wavenumber*np.pi*xf)\n",
    "    fig = plt.figure(figsize = (8, 6));\n",
    "    ax = fig.add_subplot(1,1,1);\n",
    "    if plot_sine:\n",
    "        ax.plot(xf, yf, 'r-', linewidth=2);\n",
    "    ax.plot(x,  y,  'o-', lw=2)\n",
    "\n",
    "interact(plot_sine, wavenumber=(-30,30,1), \n",
    "         grid_points=(5, 16, 1));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Try to answer the questions below with pencil and paper; then check them by modifying the code above.\n",
    "\n",
    "1. For a given number of grid points $m$, which wavenumbers $p$ will be aliased to the $p=0$ mode?  Which will be aliased to $p=1$?  Can you explain why?\n",
    "2. What is the highest frequency mode that can be represented on a given grid?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will find that the sampled frequencies (i.e., the set of values on the grid) repeat in the pattern illustrated below:\n",
    "\n",
    "![test](./aliasing_frequencies.png)\n",
    "\n",
    "For the points labeled with the same color, the values of the function when sampled on the grid are identical.  For the points with the corresponding lighter color, the values on the grid are the same except for multiplication by $-1$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aliasing as a source of numerical instability\n",
    "\n",
    "As we have seen, aliasing means that wavenumbers of magnitude greater than $\\pi m/L$ are incorrectly represented as lower wavenumbers on a grid with $m$ points.  This suggests that we shouldn't allow larger wavenumbers in our numerical solution.  For linear problems, this simply means that we should represent the initial condition by a truncated Fourier series containing modes with wavenumbers less than $\\pi m/L$.  This happens naturally when we sample the function at the grid points.  As we evolve in time, higher frequencies are not generated due to the linearity of the problem.\n",
    "\n",
    "Nonlinear problems are a different story.  Let's consider what happens when we have a quadratic term like $u^2$, as in Burgers' equation.  In general, if the grid function $u$ contains wavenumbers up to $\\pi m/L$, then $u^2$ contains frequencies up to $2 \\pi m/L$.  So each time we compute this term, we generate high frequencies that get aliased back to lower frequencies on our grid.  Clearly this has nothing to do with the correct mathematical solution and will lead to errors.  Even worse, this aliasing effect can, as it is repeated at every step, lead to an instability that causes the numerical solution to blow up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\newcommand{\\F}{\\mathcal F}\n",
    "\\newcommand{\\Finv}{{\\mathcal F}^{-1}}\n",
    "$$\n",
    "## An illustration of aliasing instability: the Korteweg-de Vries equation\n",
    "\n",
    "To see aliasing in practice, we'll consider the KdV equation, which describes certain kinds of water waves:\n",
    "\n",
    "\n",
    "$$\n",
    "u_t = -u u_x - u_{xxx}\n",
    "$$\n",
    "\n",
    "A natural pseudospectral discretization is obtained if we compute the spatial derivatives via\n",
    "\n",
    "\\begin{align}\n",
    "u_x & = \\Finv(i\\xi \\F(u)) \\\\\n",
    "u_{xxx} & = \\Finv(-i\\xi^3 \\F(u)).\n",
    "\\end{align}\n",
    "This gives\n",
    "$$\n",
    "U'(t) = -D[U] \\Finv(i\\xi \\F(U)) - \\Finv(-i\\xi^3 \\F(U)).\n",
    "$$\n",
    "\n",
    "This is identical to our discretization of Burgers' equation, except that now we have a third-derivative term.  In Fourier space, the third derivative gives a purely imaginary factor, which -- like the first derivative -- causes the solution to travel over time.  Unlike the first derivative, the third derivative term causes different wavenumber modes to travel at different speeds; this is referred to as *dispersion*.\n",
    "\n",
    "The largest-magnitude eigenvalues of the Jacobian for this semi-discretization are related to the 3rd-derivative term.  If we consider only that term, the eigenvalues are\n",
    "\n",
    "$$-i \\xi^3$$\n",
    "\n",
    "where $\\xi$ lies in the range $(-m/2, m/2)$.  So we need the time step to satisfy $\\Delta t (m/2)^3 \\in S$, where $S$ is the region of absolute stability of a given time integrator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example we'll use a 3rd-order Runge-Kutta method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk3(u,xi,rhs):\n",
    "    y2 = u + dt*rhs(u,xi)\n",
    "    y3 = 0.75*u + 0.25*(y2 + dt*rhs(y2,xi))\n",
    "    u_new = 1./3 * u + 2./3 * (y3 + dt*rhs(y3,xi))\n",
    "    return u_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the size of the imaginary axis interval contained in this method's absolute stability region:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nodepy import rk\n",
    "ssp33 = rk.loadRKM('SSP33')\n",
    "print(ssp33.imaginary_stability_interval())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll go ahead and implement our solution, making sure to set the time step according to the condition above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rhs(u, xi, equation='KdV'):\n",
    "    uhat = np.fft.fft(u)\n",
    "    if equation == 'Burgers': \n",
    "        return -u*np.real(np.fft.ifft(1j*xi*uhat)) + np.real(np.fft.ifft(-xi**2*uhat))\n",
    "    elif equation == 'KdV':\n",
    "        return -u*np.real(np.fft.ifft(1j*xi*uhat)) - np.real(np.fft.ifft(-1j*xi**3*uhat))\n",
    "    \n",
    "# Grid\n",
    "m = 256\n",
    "L = 2*np.pi\n",
    "x = np.arange(-m/2,m/2)*(L/m)\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "\n",
    "dt = 1.73/((m/2)**3)\n",
    "\n",
    "A = 25; B = 16;\n",
    "\n",
    "#u = 3*A**2/np.cosh(0.5*(A*(x+2.)))**2 + 3*B**2/np.cosh(0.5*(B*(x+1)))**2\n",
    "#tmax = 0.006\n",
    "\n",
    "# Try this one first:\n",
    "u = 1500*np.exp(-10*(x+2)**2)\n",
    "tmax = 0.005\n",
    "\n",
    "uhat2 = np.abs(np.fft.fft(u))\n",
    "\n",
    "num_plots = 50\n",
    "nplt = np.floor((tmax/num_plots)/dt)\n",
    "nmax = int(round(tmax/dt))\n",
    "\n",
    "fig = plt.figure(figsize=(12,8))\n",
    "axes = fig.add_subplot(211)\n",
    "axes2 = fig.add_subplot(212)\n",
    "line, = axes.plot(x,u,lw=3)\n",
    "line2, = axes2.semilogy(xi,uhat2)\n",
    "xi_max = np.max(np.abs(xi))\n",
    "axes2.semilogy([xi_max/2.,xi_max/2.],[1.e-6,4e8],'--r')\n",
    "axes2.semilogy([-xi_max/2.,-xi_max/2.],[1.e-6,4e8],'--r')\n",
    "axes.set_xlabel(r'$x$',fontsize=30)\n",
    "axes2.set_xlabel(r'$\\xi$',fontsize=30)\n",
    "plt.tight_layout()\n",
    "plt.close()\n",
    "\n",
    "frames = [u.copy()]\n",
    "tt = [0]\n",
    "uuhat = [uhat2]\n",
    "\n",
    "for n in range(1,nmax+1):\n",
    "    u_new = rk3(u,xi,rhs)\n",
    "\n",
    "    u = u_new.copy()\n",
    "    t = n*dt\n",
    "    # Plotting\n",
    "    if np.mod(n,nplt) == 0:\n",
    "        frames.append(u.copy())\n",
    "        tt.append(t)\n",
    "        uhat2 = np.abs(np.fft.fft(u))\n",
    "        uuhat.append(uhat2)\n",
    "        \n",
    "def plot_frame(i):\n",
    "    line.set_data(x,frames[i])\n",
    "    power_spectrum = np.abs(uuhat[i])**2\n",
    "    line2.set_data(np.sort(xi),power_spectrum[np.argsort(xi)])\n",
    "    axes.set_title('t= %.2e' % tt[i])\n",
    "    axes.set_xlim((-np.pi,np.pi))\n",
    "    axes.set_ylim((-200,3000))\n",
    "    \n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                          frames=len(frames), interval=100,\n",
    "                                          repeat=False)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the output, we're plotting the solution (top plot) and its power spectrum ($|\\hat{u}|^2$) (bottom plot).  There are a lot of interesting things to say about the solution, but for now let's focus on the Fourier transform.  Notice how the amount of energy in the high wavenumbers present (those outside the dashed red lines) remains relatively small.  Because of this, no aliasing occurs.\n",
    "\n",
    "Now change the code above to use only $m=128$ grid points.  What happens?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explanation\n",
    "\n",
    "Here we will give a somewhat simplified explanation of the blow-up just observed.  First, this blowup has nothing to do with the absoute stability condition -- when we change $m$, the time step is automatically changed in a way that will ensure absolute stability.  If you're not convinced, try taking the time step even smaller; you will still observe the blowup.\n",
    "\n",
    "By taking $m=128$, we cut by half the wavenumbers that can be represented on the grid.  As you can see from the plots, this means that after the solution begins to steepen, a significant amount of the energy present in the solution is in the upper half of the representable range of wavenumbers (i.e., outside the dashed red lines).  That means that the highest wavenumbers generated by the quadratic term will be aliased -- and they will be aliased back into that upper-half range.  This leads to a gradual and incorrect accumulation of high-wavenumber modes, easily visible in both plots.  Eventually the high-wavenumber modes dominate the numerical solution and lead to blowup.\n",
    "\n",
    "For a detailed discussion of aliasing instabilities, see Chapter 11 of John Boyd's \"Chebyshev and Fourier Spectral Methods\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filtering\n",
    "\n",
    "How can we avoid aliasing instability?  The proper approach is to ensure that the solution is well resolved, so that the instability never appears.  However, this may entail a very substantial computational cost.  One way to ensure stability even if the solution is underresolved is by *filtering*, which is also known as *dealiasing*.  In general it is unwise to rely on filtering, since it can mask the fact that the solution is not resolved (and hence  inaccurate).  But understanding filtering can give a bit more insight into aliasing instability itself."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At the most basic level, filtering means removing the modes that lead to aliasing.  This can be done by damping the high wavenumbers or simply zeroing them when computing the $(u^2)_x$ term.  The obvious approach would be to filter the upper half of all wavenumbers, but this is overkill.  In fact, it is sufficient to filter only the uppermost third.  To see why, notice that the aliased modes resulting from the lower two-thirds will appear in the uppermost third of the range of modes, and so will be filtered at the next step.\n",
    "\n",
    "A simple 2/3 filter is implemented in the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rhs(u, xi, filtr, equation='KdV'):\n",
    "    uhat = np.fft.fft(u)\n",
    "    if equation == 'Burgers': \n",
    "        return -u*np.real(np.fft.ifft(1j*xi*uhat*filtr)) \\\n",
    "                + np.real(np.fft.ifft(-xi**2*uhat))\n",
    "    elif equation == 'KdV':\n",
    "        return -u*np.real(np.fft.ifft(1j*xi*uhat*filtr)) \\\n",
    "                - np.real(np.fft.ifft(-1j*xi**3*uhat))\n",
    "    \n",
    "def rk3(u,xi,rhs,filtr):\n",
    "    y2 = u + dt*rhs(u,xi,filtr)\n",
    "    y3 = 0.75*u + 0.25*(y2 + dt*rhs(y2,xi,filtr))\n",
    "    u_new = 1./3 * u + 2./3 * (y3 + dt*rhs(y3,xi,filtr))\n",
    "    return u_new\n",
    "    \n",
    "# Grid\n",
    "m = 128\n",
    "L = 2*np.pi\n",
    "x = np.arange(-m/2,m/2)*(L/m)\n",
    "xi = np.fft.fftfreq(m)*m*2*np.pi/L\n",
    "filtr = np.ones_like(xi)\n",
    "xi_max = np.max(np.abs(xi))\n",
    "filtr[np.where(np.abs(xi)>xi_max*2./3)] = 0.\n",
    "\n",
    "dt = 1.73/((m/2)**3)\n",
    "\n",
    "A = 25; B = 16;\n",
    "u = 3*A**2/np.cosh(0.5*(A*(x+2.)))**2 + 3*B**2/np.cosh(0.5*(B*(x+1)))**2\n",
    "#u = 1500*np.exp(-10*(x+2)**2)\n",
    "tmax = 0.012\n",
    "\n",
    "uhat2 = np.abs(np.fft.fft(u))\n",
    "\n",
    "num_plots = 200\n",
    "nplt = np.floor((tmax/num_plots)/dt)\n",
    "nmax = int(round(tmax/dt))\n",
    "\n",
    "fig = plt.figure(figsize=(12,8))\n",
    "axes = fig.add_subplot(211)\n",
    "axes2 = fig.add_subplot(212)\n",
    "line, = axes.plot(x,u,lw=3)\n",
    "line2, = axes2.semilogy(xi,uhat2)\n",
    "axes2.semilogy([xi_max/2.,xi_max/2.],[1.e-6,4e8],'--r')\n",
    "axes2.semilogy([-xi_max/2.,-xi_max/2.],[1.e-6,4e8],'--r')\n",
    "axes.set_xlabel(r'$x$',fontsize=30)\n",
    "axes2.set_xlabel(r'$\\xi$',fontsize=30)\n",
    "plt.tight_layout()\n",
    "plt.close()\n",
    "\n",
    "frames = [u.copy()]\n",
    "tt = [0]\n",
    "uuhat = [uhat2]\n",
    "\n",
    "for n in range(1,nmax+1):\n",
    "    u_new = rk3(u,xi,rhs,filtr)\n",
    "\n",
    "    u = u_new.copy()\n",
    "    t = n*dt\n",
    "    # Plotting\n",
    "    if np.mod(n,nplt) == 0:\n",
    "        frames.append(u.copy())\n",
    "        tt.append(t)\n",
    "        uhat2 = np.abs(np.fft.fft(u))\n",
    "        uuhat.append(uhat2)\n",
    "        \n",
    "def plot_frame(i):\n",
    "    line.set_data(x,frames[i])\n",
    "    power_spectrum = np.abs(uuhat[i])**2\n",
    "    line2.set_data(np.sort(xi),power_spectrum[np.argsort(xi)])\n",
    "    axes.set_title('t= %.2e' % tt[i])\n",
    "    axes.set_xlim((-np.pi,np.pi))\n",
    "    axes.set_ylim((-100,3000))\n",
    "    \n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                   frames=len(frames), interval=20)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the solution remains stable, but small wiggles appear throughout the domain.  These are a hint that something is not sufficiently resolved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "The examples we have looked at behave in a way that agrees with the ideas that have been presented here.  However, the stability of nonlinear pseudospectral PDE discretizations is complicated, and it's easy to find examples that aren't fully explained by the discussion above.  You can explore this by experiments like the following.\n",
    "\n",
    "1.  Try decreasing $m$ even further in the example above.  You may find that the solution blows up for some values.  Can you prevent this by adjusting $\\Delta t$?  Can you prevent it by filtering more wavenumbers?\n",
    "\n",
    "2.  Try solving the inviscid Burgers' equation with a 2/3 filter.  For this problem, there is no value of $m$ large enough to resolve the gradients that appear, since the gradient blows up in finite time.  After the blowup time, traditionally one employs the notion of \"vanishing viscosity weak solutions\" in which a discontinuous solution is acceptable if it is given by the limit of solutions of\n",
    "$$\n",
    "u_t + uu_x = \\epsilon u_{xx}\n",
    "$$\n",
    "as $\\epsilon \\to 0$.  By using filtering, can you obtain such solutions?"
   ]
  },
  {
   "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
}