{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# p27: Solve KdV equation" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Solve the KdV equation using FFT\n", "$$\n", "u_t + u u_x + u_{xxx} = 0, \\qquad x \\in [-\\pi,\\pi]\n", "$$" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "%config InlineBackend.figure_format='svg'\n", "from matplotlib import rc\n", "rc('animation', html='jshtml')\n", "from numpy import pi,cosh,exp,round,zeros,arange,real,mod\n", "from numpy.fft import fft,ifft\n", "import matplotlib.pyplot as plt" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "# Set up grid and differentiation matrix:\n", "N = 256; dt = 0.4/N**2; x = (2*pi/N)*arange(-N/2,N/2);\n", "A, B = 25.0, 16.0\n", "u = 3*A**2/cosh(0.5*A*(x+2))**2 + 3*B**2/cosh(0.5*B*(x+1))**2\n", "v = fft(u); \n", "k = zeros(N); k[0:N//2] = arange(0,N/2); k[N//2+1:] = arange(-N/2+1,0,1)\n", "ik3 = 1j*k**3\n", "\n", "# Time-stepping by Runge-Kutta\n", "tmax = 0.006; nplt = 5 #int(round((tmax/25)/dt))\n", "nmax = int(round(tmax/dt))\n", "udata = []; udata.append(u)\n", "tdata = [0.0]\n", "for n in range(1,nmax+1):\n", " t = n*dt; g = -0.5j*dt*k\n", " E = exp(dt*ik3/2); E2 = E**2\n", " a = g * fft(real(ifft( v ))**2)\n", " b = g * fft(real(ifft( E*(v+a/2) ))**2)\n", " c = g * fft(real(ifft( E*v+b/2 ))**2)\n", " d = g * fft(real(ifft( E2*v+E*c ))**2)\n", " v = E2*v + (E2*a + 2*E*(b+c) + d)/6\n", " if mod(n,nplt) == 0:\n", " u = real(ifft(v))\n", " udata.append(u)\n", " tdata.append(t);" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "from matplotlib import animation\n", "\n", "# First set up the figure, the axis, and the plot element we want to animate\n", "fig = plt.figure()\n", "ax = plt.axes(xlim=(-pi, pi), ylim=(0, 2000))\n", "line, = ax.plot([], [], lw=2)\n", "plt.xlabel('x'); plt.ylabel('u')\n", "\n", "# initialization function: plot the background of each frame\n", "def init():\n", " line.set_data([], [])\n", " return line,\n", "\n", "# animation function. This is called sequentially\n", "def animate(i):\n", " line.set_data(x, udata[i])\n", " return line,\n", "\n", "# call the animator. blit=True means only re-draw the parts that have changed.\n", "anim = animation.FuncAnimation(fig, animate, init_func=init,\n", " frames=len(udata), interval=50, blit=True)\n", "# Save to file\n", "try:\n", " anim.save('p27.mp4', fps=20, extra_args=['-vcodec', 'libx264'])\n", "except:\n", " print(\"Cannot save mp4 file\")" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n", "\n", "\n", "\n", "\n", "

\n", " \n", "

\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "

\n", " \n", "

\n", "

\n", "\n", "\n", "\n" ], "text/plain": [ "" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Use this for inline display with controls\n", "anim\n", "\n", "# Use this for inline display of movie\n", "#from IPython.display import HTML\n", "#HTML(anim.to_html5_video())" ] } ], "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.9.6" } }, "nbformat": 4, "nbformat_minor": 1 }