{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# p06: Variable coefficient wave equation"
]
},
{
"cell_type": "code",
"execution_count": 6,
"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": 7,
"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-0.2*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": null,
"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, 2))\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",
" 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": null,
"metadata": {},
"outputs": [],
"source": [
"# Use this for inline display with controls\n",
"anim"
]
},
{
"cell_type": "code",
"execution_count": null,
"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.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 4
}