#!/usr/bin/env python
# coding: utf-8

# # Numerical solution to the wave equation in 1D and 2D
# 
# The wave equation for variable wave speed is given by (for surface waves, EM waves obey different form):
# 
# $$ \ddot{\psi} = \nabla^2\left(c^2 \psi \right) $$
# 
# 
# In 1D, the wave equation for variable wave speed is solved using a Lax-Wendroff scheme
# 
# In 2D, a Lax scheme is used instead
# 
# The numerical methods used are discussed fully in the [documentation](http://pycav.readthedocs.io/en/latest/api/pde/lax_wendroff.html)
# 
# Animations take ~ 1 min to create and display when ran in notebook
# 
# ### Importing modules and creating functions needed by PDE solver
# 
# Functions which create the initial wavepacket, its gradients along the axes and its velocity as well as the wavespeed as a function of position

# In[1]:


#NAME: Wave Equation
#DESCRIPTION: Numerically solving the wave equation in 1D and 2D.

import pycav.pde as pde
import numpy as np

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib.animation as anim

import pycav.display as display

def twoD_gaussian(XX,YY,mean,std):
    return np.exp(-((XX-mean[0])**2+(YY-mean[1])**2)/(2*std**2))

def oneD_gaussian(x,mean,std):
    return np.exp(-((x-mean)**2)/(2*std**2))

def c(x):
    return np.ones_like(x)

def velocity_1d(x,mean,std):
    def V(psi_0):
        return -c(x)*(x-mean)*oneD_gaussian(x,mean,std)/std**2
    return V

def gradient_1d(x,mean,std):
    def D(psi_0):
        return -(x-mean)*oneD_gaussian(x,mean,std)/std**2
    return D

def gradient_2d(x,y,mean,std):
    XX,YY = np.meshgrid(x,y, indexing='ij')
    def D(psi_0):
        dfdx = -(XX-mean[0])*twoD_gaussian(XX,YY,mean,std)/std**2
        dfdy = -(YY-mean[1])*twoD_gaussian(XX,YY,mean,std)/std**2
        return dfdx,dfdy
    return D

def velocity_2d(x,y,mean,std):
    XX,YY = np.meshgrid(x,y, indexing='ij')
    def V(psi_0):
        return -c_2d(x,y)*(x-mean[0])*twoD_gaussian(XX,YY,mean,std)/std**2
    return V

def c_2d(x,y):
    XX,YY = np.meshgrid(x,y, indexing='ij')
    return np.ones_like(XX)


# In[2]:


dx = 0.01
x = dx*np.arange(101)

N = 150


# ### 1D wave with reflective boundary conditions

# In[5]:


mu = 0.75
sigma = 0.05
psi_0_1d = oneD_gaussian(x,mu,sigma)

psi_1d,t = pde.LW_wave_equation(psi_0_1d,x,dx,N,c, 
            init_vel = velocity_1d(x,mu,sigma), init_grad = gradient_1d(x,mu,sigma),
            bound_cond = 'reflective')

fig1 = plt.figure(figsize = (9,6))
line = plt.plot(x,psi_1d[:,0])[0]
plt.ylim([np.min(psi_1d[:,:]),np.max(psi_1d[:,:])])

def nextframe(arg):
    line.set_data(x,psi_1d[:,arg])

animate1 = anim.FuncAnimation(fig1,nextframe, frames = N, interval = 50, repeat = False)
animate1 = display.create_animation(animate1, temp = True)
display.display_animation(animate1)


# ### 1D wave with fixed boundary conditions

# In[6]:


psi_1d,t = pde.LW_wave_equation(psi_0_1d,x,dx,N,c, 
            init_vel = velocity_1d(x,mu,sigma), init_grad = gradient_1d(x,mu,sigma),
            bound_cond = 'fixed')

fig2 = plt.figure(figsize = (9,6))
line = plt.plot(x,psi_1d[:,0])[0]
plt.ylim([np.min(psi_1d[:,:]),np.max(psi_1d[:,:])])

def nextframe(arg):
    line.set_data(x,psi_1d[:,arg])

animate2 = anim.FuncAnimation(fig2,nextframe, frames = N, interval = 50, repeat = False)
animate2 = display.create_animation(animate2, temp = True)
display.display_animation(animate2)


# ### 2D wave with reflective boundary conditions

# In[9]:


N = 200
x = dx*np.arange(101)
y = dx*np.arange(201)

XX,YY = np.meshgrid(x,y,indexing='ij')

psi_0_2d = twoD_gaussian(XX,YY,[0.5,0.8],0.1)

psi_2d,t = pde.LW_wave_equation(psi_0_2d,[x,y],dx,2*N,c_2d, a = 0.5,
                init_grad = gradient_2d(x,y,[0.5,0.8],0.1),
                bound_cond = 'reflective')

fig3 = plt.figure(figsize = (9,9))
ax = fig3.gca(projection='3d')
image = ax.plot_surface(XX.T,YY.T,psi_2d[:,:,0].T,cmap = 'plasma')

def nextframe(arg):
    ax.clear()
    ax.set_zlim3d([np.min(2*psi_2d[:,:,:]),np.max(2*psi_2d[:,:,:])])
    ax.plot_surface(XX.T,YY.T,psi_2d[:,:,2*arg].T,cmap = 'plasma')
    
animate3 = anim.FuncAnimation(fig3,nextframe, frames = int(N/2) ,interval = 50, repeat = False)
animate3 = display.create_animation(animate3, temp = True)
display.display_animation(animate3)


# ### 2D wave with reflective boundary conditions and variable wavespeed
# 
# The below simulation has the same initial conditions as the above example but the wavespeed has the function form:
# 
# $$ c(x) = \frac{1}{2}+\frac{x}{2} $$

# In[16]:


def c_2d_variable(x,y):
    XX,YY = np.meshgrid(x,y, indexing='ij')
    return 0.5+0.5*XX

psi_2d,t = pde.LW_wave_equation(psi_0_2d,[x,y],dx,2*N,c_2d_variable, a = 0.5,
                init_grad = gradient_2d(x,y,[0.5,0.8],0.1),  bound_cond = 'reflective')

fig3 = plt.figure(figsize = (9,9))
ax = fig3.gca(projection='3d')
image = ax.plot_surface(XX.T,YY.T,psi_2d[:,:,0].T,cmap = 'plasma')

def nextframe(arg):
    ax.clear()
    ax.set_zlim3d([np.min(2*psi_2d[:,:,:]),np.max(2*psi_2d[:,:,:])])
    ax.plot_surface(XX.T,YY.T,psi_2d[:,:,2*arg].T,cmap = 'plasma')
    
animate4 = anim.FuncAnimation(fig3,nextframe, frames = int(N/2),interval = 50,repeat = False)
animate4 = display.create_animation(animate4, temp = True)
display.display_animation(animate4)


# In[ ]: