#!/usr/bin/env python

from pylab import plot, show, ylim, yticks
from matplotlib.numerix import arange
from numpy import matrix, zeros
from scipy.integrate import ode

#------------------------------------
# initial conditions
#------------------------------------

neqs = 4                  # number of equations
npts = 1000               # number of points

dt = 0.1                  # step time
t1 =  dt*npts             # total time
coltab=['r','g','b','k']  # trace color

y0 = [1.0, 0.0, 5.0, 0.0] # initial values
t0 = 0.0                  # initial value 
mu = 5.0

#------------------------------------
# ode
#------------------------------------
def f(t, y, arg1): 
    return [ y[1], 
             -y[0] - arg1*y[1]*(y[0]*y[0]-1.0), 
             -y[3], 
             y[2] 
           ]

#------------------------------------

r = ode(f).set_integrator('vode', method='bdf')
r.set_initial_value(y0, t0).set_f_params(mu) 
s=matrix(zeros((npts, neqs), float))

for i in range (npts):
    if r.successful() == 0: 
        break
    s[i,:] = r.y
    r.integrate((i+1)*dt) 

#------------------------------------
# plotting
#------------------------------------
t = arange(0.0, t1, dt)
for i in range(neqs):
    plot(t, s[:,i], color=coltab[i])

show()