#!/usr/bin/env python
"""
Solve the linear convection equation

\[    u_t + a u_x = 0\]

using the following schemes:\\

\begin{itemize}
\item first order upwind (FOU)
\item Lax-Friedrich
\item Lax-Wendroff
\item Leapfrog
\item Warming and Beam
\item MacCormack
\end{itemize}

which are defined by: 

FOU:
\[u_i^{n+1} = u_i^n - s (u_i^n - u_{i-1}^n),\]

Lax-Friedrich:
\[u_i^{n+1} = 1/2 (u_{i+1}^n + u_{i-1}^n) - s/2 (u_{i+1}^n - u_{i-1}^n),\]

Lax-Wendroff:
\[u_i^{n+1} = u_i^n - s/2 (u_{i+1}^n - u_{i-1}^n) + s^2/2 (u_{i+1}^n - 2u_i^n + u_{i-1}^n),\]

Leapfrog:
\[u_i^{n+1} = u_i^n + s/2 (u_{i+1}^n - u_{i-1}^n)\]

Warming and Beam:
\[u_i^{n+1} = u_i^n - s/2 (3u_i^n - 4u_{i-1}^n + u_{i-2}^n) + s^2/2 (u_i^n - 2u_{i-1}^n + u_{i-2}^n),\]

and MacCormack's method (this is a predictor-corrector method):
\begin{eqnarray*}
ub_i^n+1 &=& u_i^n-s*(u_i^n-i_i-1^n),\\
u_i^n+1  &=& 1/2*(ub_i^n+1 + u_i^n) - s/2*(ub_i+1^n+1 - ub_i^n+1),
\end{eqnarray*}


where $s = a dt/dx$

See p.\,298 ff in Hirsch, vol.\,1.

This exercise is intended to help understand chapter 7 of Hirsch on 'Consistency, Stability, and Error Analysis'. Most examples are taken from there some have been added, some excluded.

In this exercise we don't treat flux limiters. For that, see ex9 in comp. hydrodynamics

\newpage

{\bf Tasks:}

\begin{itemize}
\item[1)] Perform a stability analysis for the FOU scheme. You should find that it is conditionally stable for $ 0 \le \sigma \le 1 $. Try out the FOU scheme for various values of sigma (which can be adjusted in the global definitions below). 

\item[2)] Repeat the same for the leapfrog scheme. You should find that it is neutrally stable: $|G| = 1$ for $|\sigma| \le 1$. 

\item[3)] Repeat the same for the Lax-Wendroff scheme. You should find that it is stable for $|\sigma| \le 1$.

\item[4)] Using CFL = 0.8 plot the results of all schemes after the same number of steps, e.g., 100 time steps. Compare.

\item[5)] Experiment with different numbers of spatial points. Compare FOU and Lax-Friedrichs for small $N \sim 100$. You will see the decoupling of the even and odd numbers which is inherent in Lax-Friedrichs but not in FOU.

\item[6)] Experiment with various combinations of values for the wave number, $k$, CFL numbers, $\sigma$, and number of mesh points per wave length. This influences the phase value, $\varphi = k \Delta x$. Compare with your stability analysis of the various schemes. You should now be able to understand why the amplitude of the wave solutions of some schemes decreases with time. What happens if you increase the number of mesh points per wavelength? Where does that put you in a plot of Diffusion error vs.\,phase angle? 

\item[7)] The group velocity of a signal is given by $v_g \doteq \frac{\dd \omega}{\dd k}$. Determine the growth factor $G$ for the leapfrog scheme. You should get
\[ G = \exp\left( - i \omega \Delta t \right) = \cos \omega \Delta t - i \sin \omega \Delta t = -i \sigma \sin \phi \pm \sqrt{1 - \sigma^2 \sin^2 \phi}.\]
Identifying the imagninary parts you obtain
\[\sin \omega \Delta t = \sigma \sin \phi,\]
from which we obtain with $\phi = k \Delta x$ and $\sigma = a \Delta t/\Delta x$
\[v_g (k) = a \frac{\cos \phi}{\cos \omega \Delta t} = a \frac{\cos \phi}{\sqrt{1 - \sigma^2 \sin^2 \phi}}.\]
Plot the group velocity and check the behaviour of the leapfrog scheme against the expected one. You should find that, indeed, for a CFL of 0.4 and $\phi = k \Delta x = \pi/4$, the Gaussian wave packet moves at about 3/4 of the real speed. This effect is entirely numerical in nature! The high frequency part of the numerical solution also propagates to the left and is reflected by the left-hand boundary, as can be seen by inserting $\phi \sim \pi$ into the above expression for the group velocity. 
\end{itemize}


Author: Bob Wimmer, wimmer@physik.uni-kiel.de

Date: December 26, 2012

"""
import numpy as np
import Gnuplot, Gnuplot.funcutils
import time


#global definitions
N    =   500     #  Number of spatial points.
n_plot= 5       #  Number of time steps to increment before updating the plot.
dx   = 1.0       #  Spatial resolution
x    = dx*np.linspace(0,N,N)        #  Spatial axis.
a = 1.           #exact group velocity/phase velocity for a linear wave
sigma = 0.4      #CFL number (generally, although not always, sigma<1 for stability. See below)
dt = sigma/a*dx  #define time step based on CFL number
print 'dt = ', dt
tstop = N/2*dt      #  Number of time steps.  

#  Direct index assignment is MUCH faster than using a spatial FOR loop, so
#  these constants are used in the update equations.  Remember that Python uses
#  zero-based indexing.
IDX1 = range(2,N-1)     #u[i]
IDX2 = range(3,N)       #u[i+1]
IDX3 = range(1,N-2)     #u[i-1]
IDX4 = range(0,N-3)     #u[i-2]
PA = 0                  #Past
PR = 1                  #Present
FU = 2                  #Future
u = np.zeros((3,N))     #one entry for past, present, and future
ub = np.zeros(N)           #u bar, used in MacCormack's method

#Define initial conditions:
#k = 1.*np.pi/200
k = np.pi/4/dx
kw = k
xn = x[range(0,N/4)]#/dx
#square wave:
#u[PR,0:N/4] = 1.
#sin wave:
#u[PR,0:N/4] = np.sin(kw*xn)
#wave packet:
sig = x[N/32]
u[PR,0:N] = np.sin(kw*x)*np.exp(-np.divide((np.subtract(x,x[N/8]))**2,(2*sig**2)))
#ramp:
#xi = np.ones(N/4)
#xu = np.divide(x[range(0,N/4)]/dx,x[N/4])
#u[PR,0:N/4] = (xi-xu)

#needed for initial conditions for the leapfrog scheme:
u[PA] = np.roll(u[PR],int(-a*dt/dx))


#analytival solution (will be propagated in plot routine):
AN = u[PR].copy()

phi = k*dx
k_min = np.pi/dx
print 'k delta x = phi =', phi, ', k = ', k, ', k_min = pi/dx = ', k_min
print 'Phi_tilde = sigma * phi = a k delta t = ', a*k*dt
epsphi = np.arctan((sigma*np.sin(phi))/(1.-2.*sigma**2*(np.sin(phi/2))**2))/(sigma*phi)
print 'a_num = a * ',epsphi, ', (a - a_num)*1500 = ',a*(1.-epsphi)*1500
G = np.sqrt(1.-4.*sigma**2*(1.-sigma**2)*(np.sin(phi/2))**4)
print 'G = ',G,', G^',N,' = ', G**N

def solver(tstop,user_action=None):

    t0 = time.clock()   #start timer to find out how much CPU time is used in solver

    t = 0.
    #print tstop*dt, dt
    while t < tstop:
        #print t, tstop*dt
#        #FOU:
#        u[FU,IDX1] = u[PR,IDX1] - sigma*(u[PR,IDX1] - u[PR,IDX3])
#        #Lax-Friedrich:
#        u[FU,IDX1] = 0.5*(u[PR,IDX2] + u[PR,IDX3]) - sigma/2*(u[PR,IDX2] - u[PR,IDX3])
#        #Lax-Wendroff:
#        u[FU,IDX1] = u[PR,IDX1] - sigma/2*(u[PR,IDX2] - u[PR,IDX3]) + sigma**2/2*(u[PR,IDX2] - 2*u[PR,IDX1] + u[PR,IDX3])
#        #Leapfrog
        u[FU,IDX1] = u[PA,IDX1] - sigma*(u[PR,IDX2] - u[PR,IDX3])
#        #Warming and Beam:
#        u[FU,IDX1] = u[PR,IDX1] - sigma/2*(3*u[PR,IDX1] - 4*u[PR,IDX3] + u[PR,IDX4]) + sigma**2/2*(u[PR,IDX1] - 2*u[PR,IDX3] + u[PR,IDX4])
#        #MacCormack:
#        ub[IDX1] = u[PR,IDX1] - sigma*(u[PR,IDX1] - u[PR,IDX3])
#        u[FU,IDX1] = 0.5*(ub[IDX1] + u[PR,IDX1]) - sigma/2*(ub[IDX2] - ub[IDX1])

        #update time step
        u[PA] = u[PR]
        u[PR] = u[FU]
        t +=dt
        if user_action is not None:
            #print u[PR]
            user_action(u[PR], x, t)       #I can do whatever I'd like here

    t1 = time.clock()
    return t1-t0, t                     #solver does NOT return solution u!
					#But it returns how long it worked (t1-t0) and the current model time

# ----------------  end of solver  -------------------------------------------



def test_solver(n_plot):
    """
    Very simple test case.
    Store the solution at every N time levels.
    Measure how long the solver actually works
    """
    g = Gnuplot.Gnuplot()#depending on your system, you may need to insert persist=1 in the parenthesis
    # Need time_level_counter as global variable since
    # it is assigned in the action function (that makes
    # a variable local to that block otherwise).
    global time_level_counter
    time_level_counter = 0
    t = 0.
    cpu = 0.
    def action(u, x, t): #this is where I can access solution u
        global time_level_counter
        if time_level_counter % n_plot == 0: #only do this every N times
            g.reset()
            g('set ylabel "y-axis [arb. units]"')   #this is how you access gnuplot commands
            g('set xlabel "x-axis [arb. units]"')
            string = 'set yrange [-1:1.5]'
            g(string)
            N0 = int(a*t/dx)
            string = 'elapsed time: ' + str(t) + ', i.e., ' + str(t/dt) + ' CFL steps'
            str2 = 'set title "' + string + '"'
            g(str2)
#            xn = x[range(0,N/4)]/dx
            #an = np.zeros(N)
            #an[0:N/4] = np.sin(k*xn)
            an = AN.copy()
            an = np.roll(an,N0)
            data_u   = Gnuplot.Data(x,u,using=(1,2), with_ = 'line', title = 'numerical')
            data_an   = Gnuplot.Data(x,an,using=(1,2), with_ = 'line', title = 'analytic')
            g.plot(data_u,data_an)
            #g.plot(data_psip)
            #z = time.sleep(0.1) 	#wait a little so you can look at it
            #print rho_res, dres
        time_level_counter += 1

    #call the solver here, note that the solution, u, is not returned
    #but cpu time is
    cpu, t = solver(tstop, user_action=action)
    print 'CPU time:', cpu, ' sec'
    print 'time: ', t
    return u
                                                   #test_solver has access to solutions through 'action'
				                   #and returns solution array to main

# ---------------------------------  end test_solver  -----------------------------


def main():
    # open output file
    #output = open("1dwave-sol.dat","w")
    u=np.zeros(N)
    u = test_solver(n_plot) 	#this is where I call the solver
    #print solutions
    #output.write(str(solutions)) 	#note: the output is an array
    #output.close()
    raw_input('Please press the return key to continue...\n')
    

main()
