# RK1D.py: Plot out time series of integration
#    steps of a 1D ODE to illustrate the fourth-order Runge-Kutta Euler method.
#
# For 1D ODE
#     dx/dt = f(x)
# the Euler method integrates this ODE by
#     x'    = x_n + f( x_n ) * Delta t
#     x_n+1 = x_n + 1/2 * [ f( x_n ) + f(x') ] * Delta t
# where Delta t is the time step and time t = n * dt.
#

# Import plotting routines
from pylab import *

# 1D ODE that has a pitchfork bifurcation
# x_dot = r * x - x * x * x
#     r in [-1,1] is the control parameter
def PitchforkODE(r,x):
	return r * x - x * x * x

# 1D Fourth-Order Runge-Kutta Integrator
def RKOneD(r,x,f,dt):
	k1 = dt * f(r,x)
	k2 = dt * f(r,x + k1/2.0)
	k3 = dt * f(r,x + k2/2.0)
	k4 = dt * f(r,x + k3)
	return x + ( k1 + 2.0 * k2 + 2.0 * k3 + k4 ) / 6.0

# Simulation parameters
# Integration time step
dt = 0.20
#
# Control parameter of the pitchfork ODE:
r = 1.0
# Set up arrays of iterates for four different initital conditions
x1 = [ 0.1]
x2 = [-0.1]
x3 = [ 2.1]
x4 = [-2.1]
# Time
t  = [ 0.0]
# The number of time steps to integrate over
N = 50

# The main loop that generates iterates and stores them
for n in xrange(0,N):
  # at each time step calculate new x(t)
  # and append to list x
  x1.append( RKOneD(r,x1[n],PitchforkODE,dt) )
  x2.append( RKOneD(r,x2[n],PitchforkODE,dt) )
  x3.append( RKOneD(r,x3[n],PitchforkODE,dt) )
  x4.append( RKOneD(r,x4[n],PitchforkODE,dt) )
  t.append( t[n] + dt )

# Setup the plot
xlabel('Time t') # set x-axis label
ylabel('x(t)') # set y-axis label
title('4th order Runge-Kutta Method: Pitchfork ODE at r= ' + str(r)) # set plot title
axis([0.0,dt*(N+1),-2.0,2.0])
# Plot the time series
plot(t,x1,'b')
plot(t,x2,'r')
plot(t,x3,'g')
plot(t,x4,'m')

# Use command below to save figure
#savefig('Euler1D', dpi=600)

# Display the plot in a window
show()