# author : gigowatts https://en.wikipedia.org/wiki/User:Gigowatts
# reference : R. J. LeVeque, "Finite difference methods for ordinary and partial differential equations", SIAM,
# {{ISBN|978-0-898716-29-0}}, chapter 10

import numpy as np
import matplotlib.pyplot as plt

L=2.5

N=500   # space discretization
a=1.0   # advection velocity
dx=L/N
dt=0.8*dx
nStep=400 # number of time steps

print 'stability constraint for Lax-Friedrichs: a*dt/dx={0}'.format(a*dt/dx)

#
# initial condition:
# - sum of two Gaussian shapes with different widths
# - Gaussian shapes are centered at x1 and x2 at t=0
#
x1=0.2
A1=2000.0
x2=0.5
A2=100.0
x  = np.arange(0.0,N)*dx
u0 = np.exp(-A1*(x-x1)**2) + np.exp(-A2*(x-x2)**2)
#u0 = np.zeros_like(x)
#u0[0.4*N:0.6*N]=1.0

#
# determine theoretical solution of advection
# at t=nStep*dt
#
x1_sol=x1+a*nStep*dt
x2_sol=x2+a*nStep*dt
u_sol=np.exp(-A1*(x-x1_sol)**2) + np.exp(-A2*(x-x2_sol)**2)

#
# Lax-Friedrichs numerical scheme
#
u=u0.copy()
iStep=0
while iStep<nStep:
    iStep +=1
    u[1:-1] = 0.5*(u[0:-2]+u[2:]) - a*dt/(2*dx)*(u[2:]-u[0:-2])
    # periodic border condition
    u[0] =u[-2]
    u[-1]=u[1]
uLF=u.copy()
    
#
# Lax-Wendroff
#
u=u0.copy()
iStep=0
while iStep<nStep:
    iStep +=1
    u[1:-1] = u[1:-1] - 0.5*a*dt/dx*(u[2:]-u[0:-2]) + 0.5*(a*dt/dx)**2*(u[0:-2]-2*u[1:-1]+u[2:])
    # periodic border condition
    u[0] =u[-2]
    u[-1]=u[1]    
uLW=u.copy()

#
# Upwind scheme
# assume a>0, so we use one-sided finite difference to approximate u_x
#
u=u0.copy()
iStep=0
while iStep<nStep:
    iStep +=1
    u[1:-1] = u[1:-1] - a*dt/dx*(u[1:-1]-u[0:-2])
    # periodic border condition
    u[0] =u[-2]
    u[-1]=u[1]
uUpwind=u.copy()

#
# leapfrog scheme
#
uOld=u0.copy()
u=u0.copy()
uNew=u0.copy()
iStep=0
while iStep<nStep:
    iStep +=1
    uNew[1:-1] = uOld[1:-1] - a*dt/dx*(u[2:]-u[0:-2])
    uOld=u.copy()
    u=uNew.copy()
    # periodic border condition
    u[0] =u[-2]
    u[-1]=u[1]
uLeapfrog=u.copy()

print 'final time is {0} after {1} time steps'.format(nStep*dt,nStep)

# #######################################
# Plot results for comparison
# #######################################

fig, axes = plt.subplots(nrows=2, ncols=2,figsize=(12,12))

axes[0,0].set_xlabel("x")
axes[0,0].set_ylabel("u(x,t=nStep*dt)")
axes[0,0].set_title("Lax-Friedrichs method")
axes[0,0].plot(x[300:500],  u_sol[300:500],'k', label="exact solution")
axes[0,0].plot(x[300:500:1],uLF[300:500:1],'r-o', markersize=5)
#axes[0,0].legend()

axes[0,1].set_xlabel("x")
axes[0,1].set_ylabel("u(x,t=nStep*dt)")
axes[0,1].set_title("Lax-Wendroff method")
axes[0,1].plot(x[300:500],  u_sol[300:500],'k', label="exact solution")
axes[0,1].plot(x[300:500:1],uLW[300:500:1],'b-o', markersize=5)
#axes[0,1].legend()

axes[1,0].set_xlabel("x")
axes[1,0].set_ylabel("u(x,t=nStep*dt)")
axes[1,0].set_title("Upwind method")
axes[1,0].plot(x[300:500],  u_sol[300:500],'k', label="exact solution")
axes[1,0].plot(x[300:500:1],uUpwind[300:500:1],'g-o', markersize=5)
#axes[1,0].legend()

axes[1,1].set_xlabel("x")
axes[1,1].set_ylabel("u(x,t=nStep*dt)")
axes[1,1].set_title("Leapfrog method")
axes[1,1].plot(x[300:500],  u_sol[300:500],'k', label="exact solution")
axes[1,1].plot(x[300:500:1],uLeapfrog[300:500:1],'c-o', markersize=5)
#axes[1,1].legend()


plt.suptitle(r"Advection equation: $\partial_t u + a \partial_x u = 0$",fontsize=15)
plt.show()
fig.savefig("advection.png",dpi=200)