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

from visual import *
from RandomArray import *

# 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

print """
Right button drag to rotate camera to view scene.
Middle button drag up or down to zoom in or out
"""
scene.title = "Lorenz differential equation"
scene.center = vector(0,0,25)
# If you can cross your eyes to see stereo, uncomment the next two lines.
#scene.stereo = 'crosseyed'
#scene.stereodepth = 2
scene.height = 600
scene.width = 600
scene.background = (0,0,0)
scene.forward = (0,1,-0.5)
scene.lights = [ vector(-0.25, -0.5, -1.0) ,vector(0.25, 0.5, 1.0) ]

# Draw axes
curve( pos = [ (-35,0,0), (35,0,0) ], color = (1,1,0), radius = 0.2 )
curve( pos = [ (0,-35,0), (0,35,0) ], color = (1,1,0), radius = 0.2 )
curve( pos = [ (0,0,0), (0,0,60) ], color = (1,1,0), radius = 0.2 )

# 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)

lorenz = curve( color = color.black, radius=0.2 )

theta = 3.0 * pi / 4.0

# 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]
	# Rotate about z to a familiar orientation
#	xtt = cos(theta) * xt - sin(theta) * yt
#	ytt = sin(theta) * xt + cos(theta) * yt
	lorenz.append( pos=(xt,yt,zt), color=(c,0.1, 1-c) )
	xState,yState,zState = xt,yt,zt

# Number of initial conditions in ensemble
nICs = 500
# Radius of ensemble
eRadius = 0.03
nIterates = 10000
State = [ ]
for i in xrange(0,nICs):
	State.append(sphere(pos=(xState+normal(0.,eRadius),\
		yState+normal(0.,eRadius),zState+normal(0.,eRadius)),
		color=(1,0.2,0.2),radius=0.2))
for n in xrange(0,nIterates):
	for i in xrange(0,nICs):
		State[i].pos[0],State[i].pos[1],State[i].pos[2] =\
			RKThreeD(sigma,R,b,State[i].pos[0],State[i].pos[1],State[i].pos[2],\
			LorenzXDot,LorenzYDot,LorenzZDot,dt)
	rate(250)