## Solution to the Lorenz Exercise

import scipy as sp
import pylab as plt
import mpl_toolkits.mplot3d
from scipy.integrate import odeint

# Constants
sigma = 10   # Prandlt number
rho = 28     # Rayleigh number
beta = 8.0/3

# Integrate!
def lorenz(n, t):
  x, y, z = n
  dxdt = sigma*(y-x)
  dydt = x*(rho-z)-y
  dzdt = x*y - beta*z
  return dxdt, dydt, dzdt
t = sp.arange(0.0, 200, 0.01)
n = odeint(lorenz, [1.0, 1.0, 1.0], t)

# Or:
#def lorenz(n, t, sigma, rho, beta):
#    ...
#t = ...
#n = odeint(lorenz, [1.0, 1.0, 1.0], t, (10, 28, 8.0/3))

# Or:
#def lorenz(n, t, sigma, rho, beta):
#    ...
#t = ...
#n = odeint(lambda n, t: lorenz(n, t, 10, 28, 8.0/3), [1.0, 1.0, 1.0], t)


x = n[:, 0]
y = n[:, 1]
z = n[:, 2]


# Plot
plt.figure().gca(projection='3d')
plt.plot(x, y, z)
plt.xlabel('x')
plt.ylabel('y')
plt.gca().set_zlabel('z')
plt.show()