# RK3D.py: Plot out time series of integration steps of a 3D ODE
#      to illustrate the fourth-order Runge-Kutta Euler method.
#
# For a 3D ODE
#     dx/dt = f(x,y)
#     dy/dt = g(x,y)
#     dz/dt = h(x,y)
# See RKThreeD() below for how the method integrates.
#

# Import plotting routines
# Use this on Linux/Unix/OS X
from pylab import *
# Use this on Windows
#from matplotlib.matlab import *

# The Lorenz 3D ODE
def LorenzXDot(sigma,R,b,x,y,z):
	return sigma * (-x + y)

def LorenzYDot(sigma,R,b,x,y,z):
	return R*x - x*z - y

def LorenzZDot(sigma,R,b,x,y,z):
	return b*z + x*y

# 3D fourth-order Runge-Kutta integrator
def RKThreeD(a,b,c,x,y,z,f,g,h,dt):
	k1x = dt * f(a,b,c,x,y,z)
	k1y = dt * g(a,b,c,x,y,z)
	k1z = dt * h(a,b,c,x,y,z)
	k2x = dt * f(a,b,c,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k2y = dt * g(a,b,c,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k2z = dt * h(a,b,c,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k3x = dt * f(a,b,c,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k3y = dt * g(a,b,c,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k3z = dt * h(a,b,c,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k4x = dt * f(a,b,c,x + k3x,y + k3y,z + k3z)
	k4y = dt * g(a,b,c,x + k3x,y + k3y,z + k3z)
	k4z = dt * h(a,b,c,x + k3x,y + k3y,z + k3z)
	x = x + ( k1x + 2.0 * k2x + 2.0 * k3x + k4x ) / 6.0
	y = y + ( k1y + 2.0 * k2y + 2.0 * k3y + k4y ) / 6.0
	z = z + ( k1z + 2.0 * k2z + 2.0 * k3z + k4z ) / 6.0
	return x,y,z

# Simulation parameters
# Integration time step
dt = 0.001
#
# Control parameters for the Lorenz ODEs:
sigma = 10.0
R     = 28.0
b     = -8.0/3.0
# Set up array of iterates and store the initial conditions
x = [5.0]
y = [5.0]
z = [5.0]
plot(x,y,'bo')
# Time
t = [0.0]
# The number of time steps to integrate over
N = 10000

# The main loop that generates the orbit, storing the states
for n in range(0,N):
  # at each time step calculate new (x,y,z)(t)
  xt,yt,zt = RKThreeD(sigma,R,b,x[n],y[n],z[n],LorenzXDot,LorenzYDot,LorenzZDot,dt)
  # and append to lists
  x.append(xt)
  y.append(yt)
  z.append(zt)

# Choose a pair of coordinates from (x,y,z) to show
# Setup the parametric plot:
xlabel('x(t)') # set x-axis label
ylabel('y(t)') # set y-axis label
# Construct plot title
PString = '(sigma,R,b) = (%g,%g,%g)' % (sigma,R,b)
title('4th order Runge-Kutta Method: Lorenz ODE at ' + PString)
axis('equal')
axis([-20.0,20.0,-20.0,20.0])
# Plot the trajectory in the phase plane
plot(x,y,'b')
axhline(0.0,color = 'k')
axvline(0.0,color = 'k')

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

# Display the plot in a window
show()