# LorenzDotSpread.py: Dot spreading on the Lorenz chaotic attractor in
#       interactive 3D.
#       Note: You can switch between Euler and Runge-Kutta integrators.

from enthought.mayavi.mlab import axes,plot3d,points3d
from numpy import arange,array
from numpy.random import normal

print """
	Rotate: Left-click & drag
	Zoom:   Right-Click & slide up/down
	Slide:  Middle-Click & drag
"""
# The Lorenz 3D ODEs
#	Original parameter values: (sigma,R,b) = (10,28,-8/3)
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 Euler integrator
def EulerThreeD(a,b,c,x,y,z,f,g,h,dt):
	return x + dt*f(a,b,c,x,y,z), y + dt*g(a,b,c,x,y,z), z + dt*h(a,b,c,x,y,z)

# 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 += ( k1x + 2.0 * k2x + 2.0 * k3x + k4x ) / 6.0
	y += ( k1y + 2.0 * k2y + 2.0 * k3y + k4y ) / 6.0
	z += ( k1z + 2.0 * k2z + 2.0 * k3z + k4z ) / 6.0
	return x,y,z

	# Simulation parameters
	# Integration time step
dt = 0.01
	#
	# Control parameters for the Lorenz ODEs:
sigma = 10.0
R	  = 28.0
b	  = -8.0/3.0
	# The number of iterations to throw away
nTransients = 200
	# The number of time steps to integrate over
nIterates = 1000

	# The main loop that generates the orbit, storing the states
xState = 5.0
yState = 5.0
zState = 5.0
for n in xrange(0,nTransients):
	xState,yState,zState = RKThreeD(sigma,R,b,xState,yState,zState,\
		LorenzXDot,LorenzYDot,LorenzZDot,dt)

  # Store the trajectory
TrajX = []
TrajY = []
TrajZ = []
Time  = []

  # Integrate the Lorenz ODEs
  #
  # Set up array of iterates and store the current state
for n in xrange(0,nIterates):
	# at each time step calculate new (x,y,z)(t)
	xt,yt,zt = RKThreeD(sigma,R,b,xState,yState,zState,\
		LorenzXDot,LorenzYDot,LorenzZDot,dt)
	# Draw lines colored by speed
	# speed = vector(xt-xState,yt-yState,zt-zState)
	# c = clip( [mag(speed) * 0.2], 0, 1 )[0]
	xState,yState,zState = xt,yt,zt

	TrajX.append(xt)
	TrajY.append(yt)
	TrajZ.append(zt)
	Time.append(n*dt)

trajectory = plot3d(TrajX,TrajY,TrajZ,Time,color=(0.,0.,1.),tube_radius=0.1)
axes(color=(0.,0.,0.),extent=[-15.0,15.0,-15.0,15.0,0.0,50.0],ranges=[-15.0,15.0,-15.0,15.0,0.0,50.0])

	# Set up ensemble
	# Number of initial conditions in ensemble
nICs = 50
	# Radius of ensemble
eRadius = 0.3
nIterates = 100
DotX = []
DotY = []
DotZ = []
for i in xrange(0,nICs):
	DotX.append(xState+normal(0.,eRadius))
	DotY.append(yState+normal(0.,eRadius))
	DotZ.append(zState+normal(0.,eRadius))

ensemble = points3d(DotX,DotY,DotZ,scale_factor=0.3,color=(1.,0.,0.))
for n in xrange(0,nIterates):
	for i in xrange(0,nICs):
		DotX[i],DotY[i],DotZ[i] =\
			RKThreeD(sigma,R,b,DotX[i],DotY[i],DotZ[i],\
			LorenzXDot,LorenzYDot,LorenzZDot,dt)
#	ensemble.mlab_source.points = array([DotX,DotY,DotZ])
	ensemble.mlab_source.points = array(map(lambda x,y,z: (x,y,z), DotX,DotY,DotZ))
#	ensemble.mlab_source.x = DotX
#	ensemble.mlab_source.y = DotY
#	ensemble.mlab_source.z = DotZ