Animate the wave function

wave.py

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation

fig = plt.figure()
ax = plt.axes(xlim=(0, np.pi), ylim=(-2, 2))
line, = ax.plot([], [], lw=2)

def init():
    line.set_data([], [])
    return line,

def f(x):
	if -np.pi/2 <= x and x <= np.pi/2:
		return x
	elif -np.pi <= x and x < -np.pi/2:
		return -(x + np.pi)
	elif np.pi/2 < x and x <= np.pi:
		return np.pi - x
	else:
		# 2Pi periodic extension
		# Most likely can make this better; was just trying to get it to work.		
		if (-2*np.pi <= x and x <= -np.pi):
			return f1(x % (2*np.pi))	
		else:
			return f1(x %(2 * np.pi) - 2*np.pi)

def animate(n):
    x = np.linspace(0, np.pi, 1000)
    y = [0.5*(f(i-n)+f(i+n)) for i in x]
    line.set_data(x, y)
    return line,

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=60, interval=200, blit=True)


# In case you just want to plot the function itself.
'''
x = np.linspace(-4*np.pi, 4*np.pi, 10000)
y = [f1(i) for i in x]
plt.plot(x,y)
'''

# Saves the animation as a .gif. Used imagemagick since I'm on a linux distribution.
anim.save('wave.gif', fps=4, writer='imagemagick')

# Generate the drawing on your desktop
#plt.show()