KdV Examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfc06113",
   "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": "code",
   "execution_count": null,
   "id": "0209d0a0",
   "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\n",
    "\n",
    "\n",
    "def rhs(u, xi, epsilon=1.0):\n",
    "    uhat = np.fft.fft(u)\n",
    "    return -u*np.real(np.fft.ifft(1j*xi*uhat)) - epsilon*np.real(np.fft.ifft(-1j*xi**3*uhat))\n",
    "    \n",
    "def solve_KdV(u0,tmax=1.,m=256,epsilon=1.0, ylims=(-100,300)):\n",
    "    \"\"\"Solve the KdV equation using Fourier spectral collocation in space\n",
    "       and SSPRK3 in time, on the domain (-pi, pi).  The input u0 should be a function.\n",
    "    \"\"\"\n",
    "    # Grid\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",
    "    u = u0(x)\n",
    "    uhat2 = np.abs(np.fft.fft(u))\n",
    "\n",
    "    num_plots = 400\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(111)\n",
    "    line, = axes.plot(x,u,lw=3)\n",
    "    xi_max = np.max(np.abs(xi))\n",
    "    axes.set_xlabel(r'$x$',fontsize=30)\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",
    "        \n",
    "    def plot_frame(i):\n",
    "        line.set_data(x,frames[i])\n",
    "        axes.set_title('t= %.2e' % tt[i])\n",
    "        axes.set_xlim((-np.pi,np.pi))\n",
    "        axes.set_ylim(ylims)\n",
    "\n",
    "    anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                       frames=len(frames), interval=100)\n",
    "\n",
    "    return HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f0cfe80",
   "metadata": {},
   "source": [
    "## Initial sinusoid\n",
    "\n",
    "Here we set up something similar to the FPUT experiment, with a single low-frequency mode as initial condition on a periodic domain.  Notice how, at some later times, the solution comes close to the initial condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bd939ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0(x):\n",
    "    return 100*np.sin(x)\n",
    "solve_KdV(u0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "784f0474",
   "metadata": {},
   "source": [
    "## Formation of a soliton train from an initial positive pulse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38c986f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0(x):\n",
    "    return 2000*np.exp(-10*(x+2)**2)\n",
    "solve_KdV(u0, tmax=0.005, ylims=(-100,3000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f191670a",
   "metadata": {},
   "source": [
    "# Interaction of two solitons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b69166",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = 25; B = 16;\n",
    "def u0(x):\n",
    "    return 3*A**2/np.cosh(0.5*(A*(x+2.)))**2 + 3*B**2/np.cosh(0.5*(B*(x+1)))**2\n",
    "solve_KdV(u0,tmax = 0.006, ylims=(-10,3000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa18aa7b",
   "metadata": {},
   "source": [
    "The next simulation shows a comparison between the propagation of a single soliton versus the interaction of two solitons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24c8e5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "u = 3*A**2/np.cosh(0.5*(A*(x+2.)))**2 + 3*B**2/np.cosh(0.5*(B*(x+1)))**2\n",
    "v = 3*A**2/np.cosh(0.5*(A*(x+2.)))**2\n",
    "\n",
    "tmax = 0.006\n",
    "\n",
    "uhat2 = np.abs(np.fft.fft(u))\n",
    "\n",
    "num_plots = 400\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(111)\n",
    "line, = axes.plot(x,u,lw=3)\n",
    "line2, = axes.plot(x,v,lw=3)\n",
    "xi_max = np.max(np.abs(xi))\n",
    "axes.set_xlabel(r'$x$',fontsize=30)\n",
    "plt.close()\n",
    "\n",
    "frames = [u.copy()]\n",
    "vframes = [v.copy()]\n",
    "tt = [0]\n",
    "uuhat = [uhat2]\n",
    "\n",
    "for n in range(1,nmax+1):\n",
    "    u_new = rk3(u,xi,rhs)\n",
    "    v_new = rk3(v,xi,rhs)\n",
    "\n",
    "    u = u_new.copy()\n",
    "    v = v_new.copy()\n",
    "    t = n*dt\n",
    "    # Plotting\n",
    "    if np.mod(n,nplt) == 0:\n",
    "        frames.append(u.copy())\n",
    "        vframes.append(v.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",
    "    line2.set_data(x,vframes[i])\n",
    "    power_spectrum = np.abs(uuhat[i])**2\n",
    "    axes.set_title('t= %.2e' % tt[i])\n",
    "    axes.set_xlim((-np.pi,np.pi))\n",
    "    axes.set_ylim((-10,3000))\n",
    "    \n",
    "anim = matplotlib.animation.FuncAnimation(fig, plot_frame,\n",
    "                                   frames=len(frames), interval=100)\n",
    "\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b44acc",
   "metadata": {},
   "source": [
    "## Formation of a dispersive shockwave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758b2d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0(x):\n",
    "    return -500*np.exp(-10*(x-2)**2)\n",
    "solve_KdV(u0, tmax=0.005, epsilon=0.1, ylims=(-600,300))"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}