"""
exampleode.py
=============

The example code helps demonstrate how to use one ODE solver
from the SCIPY package: scipy.integrate.odeint()
There is a more object oriented solver in SCIPY called ode
which may have some features that odeint does not.

The example set of equations is the van der Pol oscillator:
         y''-mu*(1-y^2)*y'+y=0    y(0)=2,  y'(0)=0
The first step is to rewrite as a system of equations:
         y1'=y2      y2'=mu*(1-y^2)*y'+y    y1(0)=2,  y2(0)=0
The function vanderpol() returns the RHS vector for this system.

The function run_vanderpol() sets up parameters, runs the solver,
and returns the solution. 
The function draw_figures() plots a few figures with this data.
"""
from numpy import *
from scipy.integrate import odeint
import matplotlib.pyplot as plt

def vanderpol(y,t,mu):
    """ Return the derivative vector for the van der Pol equations."""
    y1= y[0]
    y2= y[1]
    dy1=y2
    dy2=mu*(1-y1**2)*y2-y1
    return [dy1, dy2]

def run_vanderpol(yinit=[2,0], tfinal=20, mu=2):
    """ Example for how to run odeint.

    More info found in the doc_string. In ipython type odeint?
    """
    times = linspace(0,tfinal,200)

    rtol=1e-6
    atol=1e-10

    y = odeint(vanderpol, yinit, times, args= (mu,), rtol=rtol, atol=atol)
    return y,times

def draw_figure(y,t):
    plt.subplot(2,2,1)
    plt.plot(t, y[:,0],'-ob')
    plt.title("Time Profile of y1")
    plt.subplot(2,2,3)
    plt.plot(t, y[:,1],'-og')
    plt.title("Time Profile of y2")
    # plt.subplot2grid((2,2),(0,1),rowspan=2)
    # in 2x2 grid of subplots, draw in (0,1)-->upper right, 
    # and make the plot span two rows of plots (height is twice normal)
    plt.subplot2grid((2,2),(0,1),rowspan=2)
    plt.plot(y[:,0], y[:,1],'-ok')
    plt.title("Phase Portrait")

if __name__ == "__main__":
    y,t = run_vanderpol([2,0], tfinal=20, mu=2)
    draw_figure(y,t)
    plt.show()