# poincare.py
'''
Create a low dimensional slice of the Lorenz attractor known as a Poincare Section, Return map, or Recurrence map.

We look for z points that are local maximums and plot them versus the previous local maximum.
'''

import math
import matplotlib.pyplot as plt


DT = 0.01
N = 100000

def lorenz(xyz, s=10, p=28, b=8/3):
  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)


currs = []
prevs = []
xyz = [(0,0,0) for x in range(N+1)]
xyz[0] = (0, 1.0, 1.05)

localMax = 0   # dummy start value
for i in range(N):
  xyz[i+1] = tuple(map(lambda x,f: x+f*DT, 
                       xyz[i], lorenz(xyz[i])))
  if i > 100:   # let things settle down
    zOld = xyz[i-1][2]
    z = xyz[i][2]
    zNew = xyz[i+1][2]
    if zNew < z and zOld < z:
      prevMax = localMax
      localMax = z
      currs.append(localMax)
      prevs.append(prevMax)

plt.scatter(prevs[1:], currs[1:], s=5)  # s is point size
            # skip dummy value
plt.xlabel("Zmax Prev")
plt.ylabel("Zmax Curr")
plt.title("Poincare section for Lorenz")
plt.show()