######################
# Physics 730: PDEs 1: Solution to the 1D Waveequation w/ animation
# JRG, UM Physics
# NOTE: This needs to be run with the QT4 graphics library:
# Start ipython like: ipython --matplotlib=qt
######################

from pylab import *
import time
import matplotlib.animation as animation

def solver0(I,f,c,L,n,dx,tstop,dt=0.0,method='v'):
    
    #Set up an empty list container for solutions at ALL time steps
    fullU=[]
    
    #if dt is not specified, make largest that is stable
    	
    if dt <= 0: dt=dx/max(c(x))
    
    C2=(dt/dx)**2
    dt2=dt*dt
    
    up=zeros(n+1)		#solution array
    u=up.copy()			# solution at t-dt
    um=up.copy()		#solution at t-2dt

    t=0.0
    
    #Set up intitial conditions
    for i in range(n+1):
        u[i]=I(x[i])
    for i in range(1,n):
        um[i]=u[i] + 0.5*C2*(u[i-1] - 2*u[i] + u[i+1]) + dt2*f(x[i],t)
    
    #Set up Dirichlet boundary conditions
    um[0]=0.
    um[n]=0.
    
    while t <= tstop:
        t_old = t; t+=dt
        if method == 's':
            for i in range(1,n):
                up[i] = -um[i] +2*u[i] + c(u[i])**2*C2*(u[i-1] - 2*u[i] + u[i+1]) + dt2*f(x[i],t_old)
        elif method == 'v':
            up[1:n] = -um[1:n] +2*u[1:n] +c(u[1:n])**2*C2*(u[0:n-1] - 2*u[1:n] +u[2:n+1]) + dt2*f(x[1:n],t_old)
        else:
            print "Method needs to be 's' for scalar or 'v' for vector.  \n Exiting ..."
            break
        up[0] = 0.; up[n] = 0.
        fullU.append(u)

        um=u.copy(); u=up.copy()
    return dt, fullU, method

def animate(fullU):
	animFig=figure()
	line,=plot(x,fullU[0],'b-')
	cdata=c(x)/c(x).max()
	line2,=plot(x,cdata,'g--')
	speed=10*dt
	
	def init():
		line.set_ydata(zeros(len(x))+np.nan)
		line2.set_ydata(cdata)
		ylim([-1,1])
		xlabel('X',fontsize=18)
		ylabel('$\psi(x)$',fontsize=20)
		title('Solution to the 1D Wave Equation')
		return line, line2,
		
	def update(i): 
		line.set_ydata(fullU[i])
		#line2.set_ydata(c(x))
		time=float(i)*dt
		#text(max(x)*0.9,max(fullU[0])*0.9,'Time=%2.2f s'%time)
		ax=axes()
		# Adding the time on the plot is nice, but really bogs down animation
		#txt=text(0.8,0.9,'Time=%2.2f s'%time,transform = ax.transAxes)
		return line, line2
		
	ani = animation.FuncAnimation(
		animFig,
			update,
			range(len(fullU)),
			blit=True,
			interval=speed,
			init_func=init
			)
	show()
	
def waterfallPlot(fullU,dstep=25,offsetFactor=0.5):
	numPlots=0
	offset=0.0
	for i in range(len(fullU)):
		if offset<max(fullU[i]): offset=offsetFactor*max(fullU[i])
		
	for step in range(0,len(fullU),dstep):
		time=float(step)*dt
		plot(x,fullU[step]+numPlots*offset,label='t=%2.1f'%time)
		numPlots+=1
	
	xlabel('X',fontsize=18)
	ylabel('$\psi(x)$',fontsize=20)
	if numPlots<10: legend()
	show()
	
#### Main part of program
# Set up initial condition (I) and a driver function (f)
def I(x): return   exp(-5*(x-L/3.)**2) #sin(2*pi*x/L)
def f(x,t): return  0.0 #6*sin(2*pi*0.5*t+pi/4)*exp(-2.0*((x)*3)**2) 

### Function for the wave speed

def c(x,profile='constant'):
	# This make a step function at posn (as % of domain)
	if profile=='step':
		posn=0.5
		allC=zeros(x.size)
		lowX=x[x<x[-1]*posn]
		allC[0:lowX.size]=1.0
		allC[lowX.size:x.size]=2.0
		return allC
	
	# This returns a smoothly varying wavespeed over domain
	if profile=='smooth': return 1+0.05*x**2

	# This just returns a constant wavespeed over the whole domain
	if profile=='constant': return zeros(x.size)+2.0

def cNL(u,a=1.0,c0=1.0):
	return c0 +a*abs(u)

n=200; tstop=30; L=10.
dx=L/float(n)
x=linspace(0,L,n+1)



t0=time.time()

print "Computing Solution..."
dt,fullU,method=solver0(I,f,c,L,n,dx,tstop,method='v')
tf=time.time()

print "Total time for method %s = %2.5f seconds." %(method,tf-t0) 


#waterfallPlot(fullU,offsetFactor=0.2, dstep=10)
animate(fullU)