############
############ This program creates a movie of a function of time and position and saves it in movie.gif
############
############ Written by Paul Batzing as an example for FYS2140 students.
############

############
############ DO NOT CHANGE
############
from pylab import * #to plot and for math
from matplotlib import pyplot
from matplotlib import animation
from subprocess import call #to be able to delete temporary files
############
############ END OF DO NOT CHANGE
############



############
############ Replace with the function you want to plot.
############ This time it is a function of time and space (some kind of wave packet).
############
############ The function you want to plot
def f(x,t):
    ## You can break up the calculation in bits.
    ft = sin(0.65*x-1.19*t)
    fx = sin (0.1*x-0.65*0.01/1.19*t)
    f = 2*fx*ft # OBS! THIS IS AN EXAMPLE, REPLACE WITH THE FUNCTION YOU WANT TO PLOT.
    return f # This returns the value you defined in f
############
############ End of function
############


############
############ Replace with the range you want to plot.
############ Now we need both a position and a time range.
############
############ Define the ranges you want to plot over:
xmin = 0 #m
xmax = 150 #m
tmin = 0 #s
tmax = 100 #s
#t = linspace(tmin, tmax,100)
############
############ End of range.
############

############
############ Setup of the figure.
############
fig = pyplot.figure()
ax = pyplot.axes(xlim=(xmin,xmax),ylim=(-2,2))
line, = ax.plot([],[],lw=2)
############
############ Setup of the figure.
############

############
############ Initialize the plot.
############
def init():
    line.set_data([], [])
    return line,
############
############ End of the initialization.
############

############
############ Define the animator.
############
def animate(t):
    x = linspace(xmin, xmax,1000)
    y = f(x,t)
    line.set_data(x,y)
    return line,
############
############ End of the animator.
############

############
############ DO NOT CHANGE!
############ Animate the function and show the result.
anim = animation.FuncAnimation(fig, animate, init_func=init,frames=200, interval=tmax, blit=True)
pyplot.show()
############
############ END OF DO NOT CHANGE!
############