{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Numerical solution to the wave equation in 1D and 2D\n",
"\n",
"The wave equation for variable wave speed is given by (for surface waves, EM waves obey different form):\n",
"\n",
"$$ \\ddot{\\psi} = \\nabla^2\\left(c^2 \\psi \\right) $$\n",
"\n",
"\n",
"In 1D, the wave equation for variable wave speed is solved using a Lax-Wendroff scheme\n",
"\n",
"In 2D, a Lax scheme is used instead\n",
"\n",
"The numerical methods used are discussed fully in the [documentation](http://pycav.readthedocs.io/en/latest/api/pde/lax_wendroff.html)\n",
"\n",
"Animations take ~ 1 min to create and display when ran in notebook\n",
"\n",
"### Importing modules and creating functions needed by PDE solver\n",
"\n",
"Functions which create the initial wavepacket, its gradients along the axes and its velocity as well as the wavespeed as a function of position"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"#NAME: Wave Equation\n",
"#DESCRIPTION: Numerically solving the wave equation in 1D and 2D.\n",
"\n",
"import pycav.pde as pde\n",
"import numpy as np\n",
"\n",
"from mpl_toolkits.mplot3d import Axes3D\n",
"import matplotlib.pyplot as plt\n",
"import matplotlib.animation as anim\n",
"\n",
"import pycav.display as display\n",
"\n",
"def twoD_gaussian(XX,YY,mean,std):\n",
" return np.exp(-((XX-mean[0])**2+(YY-mean[1])**2)/(2*std**2))\n",
"\n",
"def oneD_gaussian(x,mean,std):\n",
" return np.exp(-((x-mean)**2)/(2*std**2))\n",
"\n",
"def c(x):\n",
" return np.ones_like(x)\n",
"\n",
"def velocity_1d(x,mean,std):\n",
" def V(psi_0):\n",
" return -c(x)*(x-mean)*oneD_gaussian(x,mean,std)/std**2\n",
" return V\n",
"\n",
"def gradient_1d(x,mean,std):\n",
" def D(psi_0):\n",
" return -(x-mean)*oneD_gaussian(x,mean,std)/std**2\n",
" return D\n",
"\n",
"def gradient_2d(x,y,mean,std):\n",
" XX,YY = np.meshgrid(x,y, indexing='ij')\n",
" def D(psi_0):\n",
" dfdx = -(XX-mean[0])*twoD_gaussian(XX,YY,mean,std)/std**2\n",
" dfdy = -(YY-mean[1])*twoD_gaussian(XX,YY,mean,std)/std**2\n",
" return dfdx,dfdy\n",
" return D\n",
"\n",
"def velocity_2d(x,y,mean,std):\n",
" XX,YY = np.meshgrid(x,y, indexing='ij')\n",
" def V(psi_0):\n",
" return -c_2d(x,y)*(x-mean[0])*twoD_gaussian(XX,YY,mean,std)/std**2\n",
" return V\n",
"\n",
"def c_2d(x,y):\n",
" XX,YY = np.meshgrid(x,y, indexing='ij')\n",
" return np.ones_like(XX)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"dx = 0.01\n",
"x = dx*np.arange(101)\n",
"\n",
"N = 150"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 1D wave with reflective boundary conditions"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"mu = 0.75\n",
"sigma = 0.05\n",
"psi_0_1d = oneD_gaussian(x,mu,sigma)\n",
"\n",
"psi_1d,t = pde.LW_wave_equation(psi_0_1d,x,dx,N,c, \n",
" init_vel = velocity_1d(x,mu,sigma), init_grad = gradient_1d(x,mu,sigma),\n",
" bound_cond = 'reflective')\n",
"\n",
"fig1 = plt.figure(figsize = (9,6))\n",
"line = plt.plot(x,psi_1d[:,0])[0]\n",
"plt.ylim([np.min(psi_1d[:,:]),np.max(psi_1d[:,:])])\n",
"\n",
"def nextframe(arg):\n",
" line.set_data(x,psi_1d[:,arg])\n",
"\n",
"animate1 = anim.FuncAnimation(fig1,nextframe, frames = N, interval = 50, repeat = False)\n",
"animate1 = display.create_animation(animate1, temp = True)\n",
"display.display_animation(animate1)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 1D wave with fixed boundary conditions"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"psi_1d,t = pde.LW_wave_equation(psi_0_1d,x,dx,N,c, \n",
" init_vel = velocity_1d(x,mu,sigma), init_grad = gradient_1d(x,mu,sigma),\n",
" bound_cond = 'fixed')\n",
"\n",
"fig2 = plt.figure(figsize = (9,6))\n",
"line = plt.plot(x,psi_1d[:,0])[0]\n",
"plt.ylim([np.min(psi_1d[:,:]),np.max(psi_1d[:,:])])\n",
"\n",
"def nextframe(arg):\n",
" line.set_data(x,psi_1d[:,arg])\n",
"\n",
"animate2 = anim.FuncAnimation(fig2,nextframe, frames = N, interval = 50, repeat = False)\n",
"animate2 = display.create_animation(animate2, temp = True)\n",
"display.display_animation(animate2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2D wave with reflective boundary conditions"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"N = 200\n",
"x = dx*np.arange(101)\n",
"y = dx*np.arange(201)\n",
"\n",
"XX,YY = np.meshgrid(x,y,indexing='ij')\n",
"\n",
"psi_0_2d = twoD_gaussian(XX,YY,[0.5,0.8],0.1)\n",
"\n",
"psi_2d,t = pde.LW_wave_equation(psi_0_2d,[x,y],dx,2*N,c_2d, a = 0.5,\n",
" init_grad = gradient_2d(x,y,[0.5,0.8],0.1),\n",
" bound_cond = 'reflective')\n",
"\n",
"fig3 = plt.figure(figsize = (9,9))\n",
"ax = fig3.gca(projection='3d')\n",
"image = ax.plot_surface(XX.T,YY.T,psi_2d[:,:,0].T,cmap = 'plasma')\n",
"\n",
"def nextframe(arg):\n",
" ax.clear()\n",
" ax.set_zlim3d([np.min(2*psi_2d[:,:,:]),np.max(2*psi_2d[:,:,:])])\n",
" ax.plot_surface(XX.T,YY.T,psi_2d[:,:,2*arg].T,cmap = 'plasma')\n",
" \n",
"animate3 = anim.FuncAnimation(fig3,nextframe, frames = int(N/2) ,interval = 50, repeat = False)\n",
"animate3 = display.create_animation(animate3, temp = True)\n",
"display.display_animation(animate3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2D wave with reflective boundary conditions and variable wavespeed\n",
"\n",
"The below simulation has the same initial conditions as the above example but the wavespeed has the function form:\n",
"\n",
"$$ c(x) = \\frac{1}{2}+\\frac{x}{2} $$"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
"