{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# p06: Variable coefficient wave equation\n", "\n", "Solve using FFT\n", "\n", "$$\n", "u_t + c(x) u_x = 0, \\qquad x \\in (0,2\\pi), \\qquad c(x) = \\frac{1}{5} + \\sin^2(x-1)\n", "$$\n", "\n", "with initial condition\n", "\n", "$$\n", "u(x,0) = u_0(x) = \\exp(-100(x-1)^2)\n", "$$\n", "\n", "periodic boundary conditions and Leapfrog scheme in time: for $n=1,2,\\ldots$\n", "\n", "$$\n", "v_j^{n+1} = v_j^{n-1} - (2 \\Delta t) c(x_j) w_j^n, \\qquad w_j^n \\approx u_x(x_j, t_n)\n", "$$\n", "\n", "Take initial values\n", "\n", "$$\n", "v_j^0 = u_0(x_j), \\qquad v_j^{-1} = u_0(x_j + c(x_j)\\Delta t)\n", "$$" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_format='svg'\n", "from matplotlib.collections import PolyCollection\n", "from numpy import pi,linspace,sin,cos,exp,round,zeros,arange,real\n", "from numpy.fft import fft,ifft\n", "from matplotlib.pyplot import figure" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n" ], "text/plain": [ "

" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# Set up grid and differentiation matrix:\n", "N = 128; h = 2*pi/N; x = h*arange(1,N+1);\n", "t = 0.0; dt = h/4.0\n", "tmax = 8.0; tplot = 0.15;\n", "plotgap = int(round(tplot/dt)); dt = tplot/plotgap;\n", "nplots = int(round(tmax/tplot))\n", "\n", "c = 0.2 + sin(x-1.0)**2\n", "v = exp(-100.0*(x-1.0)**2); vold = exp(-100.0*(x+c*dt-1.0)**2);\n", "\n", "# wave numbers\n", "ik = 1j*zeros(N)\n", "ik[0:N//2+1] = 1j*arange(0,N//2+1)\n", "ik[N//2+1:] = 1j*arange(-N//2+1,0,1)\n", "\n", "# Time-stepping by leap-frog formula\n", "data = []; data.append(list(zip(x, v)))\n", "tdata = []; tdata.append(0.0)\n", "for i in range(1,nplots):\n", " for n in range(plotgap):\n", " t = t + dt\n", " v_hat = fft(v)\n", " w_hat = ik * v_hat; w_hat[N//2] = 0.0\n", " w = real(ifft(w_hat))\n", " vnew = vold - 2.0*dt*c*w\n", " vold = v; v = vnew;\n", " data.append(list(zip(x, v)))\n", " tdata.append(t);\n", "\n", "fig = figure(figsize=(12,10))\n", "ax = fig.add_subplot(111,projection='3d')\n", "poly = PolyCollection(data, closed=False, facecolors='white', edgecolors='red')\n", "poly.set_alpha(1)\n", "ax.add_collection3d(poly, zs=tdata, zdir='y')\n", "ax.set_xlabel('x'); ax.set_ylabel('t'); ax.set_zlabel('v')\n", "ax.set_xlim3d(0, 2*pi); ax.set_ylim3d(0, 8); ax.set_zlim3d(0, 4)\n", "ax.view_init(40,-70)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Make animation" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "from matplotlib import animation\n", "from matplotlib import rc\n", "rc('animation', html='jshtml')\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=(0, 2*pi), ylim=(-0.1, 1.1))\n", "line, = ax.plot([], [], 'r-', lw=2, label='u')\n", "ax2 = ax.twinx()\n", "speed, = ax2.plot(x,c,'k:',lw=1,label='c')\n", "ax.legend(loc='upper left'), ax2.legend(loc='upper right')\n", "ax.set_xlabel('x'); ax.set_ylabel('u'); ax2.set_ylabel('c')\n", "plt.close();\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", " x, v = zip(*data[i])\n", " line.set_data(x, v)\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(data), interval=50, blit=True)\n", "# Save to file\n", "try:\n", " anim.save('p06.mp4', fps=20, extra_args=['-vcodec', 'libx264'])\n", "except:\n", " print(\"Cannot save mp4 file\")" ] }, { "cell_type": "code", "execution_count": 26, "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": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Use this for inline display with controls\n", "anim" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The dotted line shows the speed function $c(x)$." ] }, { "cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [], "source": [ "# Use this for inline display of movie\n", "#from IPython.display import HTML\n", "#HTML(anim.to_html5_video())" ] } ], "metadata": { "anaconda-cloud": {}, "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.12.8" } }, "nbformat": 4, "nbformat_minor": 4 }