###################### # 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