# lorenz.py
"""
  Plot of Edward Lorenz's 1963 "Deterministic 
  Nonperiodic Flow" in 3D space
  https://en.wikipedia.org/wiki/Lorenz_system
"""

import matplotlib.pyplot as plt
from frange import *

DT = 0.01
N = 10000

def lorenz(xyz, s=10, p=28, b=8/3):
  """
  xyz : Point in 3D space as a tuple
  s, p, b : Parameters defining the Lorenz attractor
  Returns Lorenz attractor's partial derivatives at xyz
  """
  x, y, z = xyz
  x_dot = s*(y - x)
  y_dot = p*x - y - x*z
  z_dot = x*y - b*z
  return (x_dot, y_dot, z_dot)

xyz = [(0,0,0) for x in range(N+1)]
xyz[0] = (0, 1.0, 1.05)
for i in range(N):
  xyz[i+1] = tuple(map(lambda x,f: x+f*DT, 
                                xyz[i], lorenz(xyz[i])))
                # xyz[i] + lorenz(xyz[i]) * DT
fig = plt.figure(figsize=(7, 7))
ax = fig.add_subplot(projection ='3d')
xs, ys, zs = zip(*xyz)
ax.plot(xs, ys, zs, lw=0.5)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_title("Lorenz Attractor")
plt.show()