Oscillator animation

oscillatorAnimation.py

# Animation of harmonic oscillator, with Verlet integration

from __future__ import division
import numpy as np
from numpy import sin, cos
from scipy.constants import c,e,h,hbar,u,pi
import matplotlib.animation as animation

delta = 1e-2                # spatial discretization
x = np.arange(-2,2,delta)

fig, ax = plt.subplots()

# ***** ANIMATION ***********
def V(x): return 0.5*x**2  # m=1, omega=1
def a(x,d=delta): return -( V(x+d) - V(x) )/d
def energy(x,v): return 0.5*v**2 + V(x)

line, = ax.plot(x,V(x),lw=2)
x0,v0 = np.random.uniform(-1,1,size=2)
blob, = ax.plot(x0,V(x0),'ro',ms=16)

xi,vi = x0,v0

text_template = 'energy = %.2g'
 
text = ax.text(0.25,0.5,text_template%(energy(x0,v0)), transform=ax.transAxes, family='monospace',fontsize=20,weight='heavy') 

def animate(i,h=1e-1):
    global xi, vi, a
    t = i*h
    ## Verlet algorithm
    xold = xi
    xi += vi*h + 0.5*a(xi)*h**2
    vi += 0.5*( a(xold) + a(xi) )*h
    
    blob.set_xdata( xi )
    blob.set_ydata( V(xi) )  # update the data
    text.set_text( text_template%(energy(xi,vi)) )
    return (blob,) + (text,)

def init(): #Init only required for blitting to give a clean slate.
    blob.set_ydata(np.ma.array(x, mask=True))
    text.set_text(" ")
    return (blob,) + (text,)

ani = animation.FuncAnimation(fig, animate, range(1000), init_func=init,
    interval=0.2, blit=True)
plt.show()