from __future__ import division
from pylab import *

def euler(f, g, q0, p0, h, N):
    q = zeros((N+1, size(q0)))
    p = zeros((N+1, size(p0)))
    q[0] = q0
    p[0] = p0
    for n in range(N):
        q[n+1] = q[n] + h*f(p[n])
        p[n+1] = p[n] + h*g(q[n])
    return q, p

def shoPhasePlot(method):
    figure()

    #create contour plot of Hamiltonian H, which shows the exact solutions
    Q, P = meshgrid(linspace(-pi,pi,256),linspace(-pi,pi,256))
    H = 0.5*(P**2 + Q**2)
    contour(Q, P, H, 0.5*linspace(-pi,pi,16)**2, colors='r', linestyles='-')
    
    def f(p):
        return p
    
    def g(q):
        return -q

    #plot numerical solutions
    for p0 in linspace(-pi,pi,16):
        q, p = method(f, g, 0, p0, 0.25, 400)
        plot(q, p, color='b', linestyle='', marker='.')

    title("SHO phase plot for " + method.__name__)

def pendulumPhasePlot(method):
    figure()

    #create contour plot of Hamiltonian H, which shows the exact solutions
    Q, P = meshgrid(linspace(-pi,pi,256),linspace(-pi,pi,256))
    H = 0.5*P**2 - cos(Q)
    contour(Q, P, H, 0.5*linspace(-3,3,16)**2 - 1, colors='r', linestyles='-')

    def f(p):
        return p
    
    def g(q):
        return -sin(q)
    
    for p0 in linspace(-3,3,16):
        q, p = method(f, g, 0, p0, 0.25, 400)
        plot(mod(q+pi,2*pi)-pi, p, color='b', linestyle='', marker='.')

    title("Pendulum phase plot for " + method.__name__)