%matplotlib inline
from pylab import *
from numpy import *

#dimensions of the computational domain:
maxx = 10.
maxt = 20.

#flow velocity
v = 0.5

#discretization parameters
nx = 100   # number of unknown grid points in spatial direction
CFL = 1.0  # Courant constant v*dt/dx

def wave_init(maxx, maxt, v, nx, CFL):
    #choose time step according to CFL condition
    dx = maxx/(nx+1)
    dt = CFL*dx/v
    nt = int(maxt/dt)+1
    
    #define array for storing the solution
    U = zeros((nt, nx+2))
    
    x = arange(nx+2)*dx
    t = arange(nt)*dt
    return U, dx, dt, x, t

def convection_solve(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        U[it+1,xint] = U[it,xint] - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint-1])

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
#U[0,:] = sin(x*5)
plot(U[0,:])
ylim(0,1.1)

convection_solve(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_Lax(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        U[it+1,xint] = 0.5*(U[it,xint+1]+U[it,xint-1]) - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint-1])

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_Lax(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=1.05)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_Lax(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_Lax(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_upwind(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        U[it+1,xint] = U[it,xint] - v*dt/dx*(U[it,xint] - U[it, xint-1])

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_upwind(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LeapFrog(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    #start the solution with upwind method
    U[1,xint] = U[0,xint] - v*dt/dx*(U[0,xint] - U[0, xint-1])
    for it in range(1, nt-1):
        U[it+1,xint] = U[it-1,xint] - v*dt/dx*(U[it,xint+1] - U[it, xint-1])

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LeapFrog(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LeapFrog(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    #start the solution with upwind method
    U[1,xint] = U[0,xint] - v*dt/dx*(U[0,xint] - U[0, xint-1])
    for it in range(1, nt-1):
        U[it+1,xint] = U[it-1,xint] - v*dt/dx*(U[it,xint+1] - U[it, xint-1])
                
        #extrapolate values on the outflow boundary
        U[it+1,-1] = U[it+1, -2]
U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LeapFrog(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LeapFrog(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    #start and smooth the solution with upwind method
    for it in range(0,min(50, nt-1)):
        U[it+1,xint] = U[it,xint] - v*dt/dx*(U[it,xint] - U[it, xint-1])
        #U[it+1,xint] = 0.5*(U[it,xint+1]+U[it,xint-1]) - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint-1])
    for it in range(min(50, nt-1),nt-1):
        U[it+1,xint] = U[it-1,xint] - v*dt/dx*(U[it,xint+1] - U[it, xint-1])
                
        #extrapolate values on the outflow boundary
        U[it+1,-1] = U[it+1, -2]*2 - U[it+1, -3]
U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LeapFrog(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LaxWendroff(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        # less efficient but more readable:
        Uplus  = 0.5*(U[it,xint+1] + U[it, xint]) - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint])
        Uminus = 0.5*(U[it,xint] + U[it, xint-1]) - 0.5*v*dt/dx*(U[it,xint] - U[it, xint-1])
        U[it+1,xint] = U[it,xint] - v*dt/dx*(Uplus - Uminus)

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LaxWendroff(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LaxWendroff(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        # less efficient but more readable:
        Uplus  = 0.5*(U[it,xint+1] + U[it, xint]) - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint])
        Uminus = 0.5*(U[it,xint] + U[it, xint-1]) - 0.5*v*dt/dx*(U[it,xint] - U[it, xint-1])
        U[it+1,xint] = U[it,xint] - v*dt/dx*(Uplus - Uminus)
        
        #extrapolate values on the outflow boundary
        U[it+1,-1] = U[it+1, -2]*2 - U[it+1,-3]

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LaxWendroff(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

def convection_solve_LaxWendroff(U, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,min(50, nt-1)):
        U[it+1,xint] = U[it,xint] - v*dt/dx*(U[it,xint] - U[it, xint-1])
    for it in range(min(50, nt-1),nt-1):
        # less efficient but more readable:
        Uplus  = 0.5*(U[it,xint+1] + U[it, xint]) - 0.5*v*dt/dx*(U[it,xint+1] - U[it, xint])
        Uminus = 0.5*(U[it,xint] + U[it, xint-1]) - 0.5*v*dt/dx*(U[it,xint] - U[it, xint-1])
        U[it+1,xint] = U[it,xint] - v*dt/dx*(Uplus - Uminus)
        
        #extrapolate values on the outflow boundary
        U[it+1,-1] = U[it+1, -2]*2 - U[it+1,-3]

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=.1)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_LaxWendroff(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)

from scipy.sparse import dia_matrix, eye
from scipy.sparse.linalg import factorized
def d1matrix(nelem):
    elements = ones((3,nelem))
    elements[1,:] *= 0
    elements[0,:] *= -1
    return dia_matrix((elements, [-1,0,1]), shape=(nelem,nelem)).tocsc()

def convection_solve_CN(U, dx, dt, nx, nt):
    alpha = -v*dt/(dx*4.)
    M1 = eye(nx)-d1matrix(nx)*alpha
    M2 = eye(nx)+d1matrix(nx)*alpha
    LU = factorized(M1.tocsc())
    for it in range(0,nt-1):
        U[it+1,1:-1] = LU(M2.dot(U[it,1:-1]))
        
        #extrapolate values on the outflow boundary
        U[it+1,-1] = U[it+1, -2]

        

U, dx, dt, x, t = wave_init(maxx, maxt, v, nx, CFL=2)
U[0,:] = 0.
U[0, (x<maxx*0.4) & (x>maxx*0.2)] = 1.
convection_solve_CN(U, dx, dt, nx, len(t))
pcolormesh(U, rasterized=True, vmin=-2, vmax=2)