"""
================
Lorenz attractor
================
Este é um exeplo do Atrator de Edward Lorenz's (1963) "Deterministic Nonperiodic Flow":
 https://journals.ametsoc.org/view/journals/atsc/20/2/1520-0469_1963_020_0130_dnf_2_0_co_2.xml
"""

from numpy import array, empty, arange, sqrt
import matplotlib.pyplot as plt
from matplotlib import ticker
plt.rcParams['text.usetex'] = True

def lorenz(xyz, *, s=10, r=28, b=8/3):
    """
    Parâmetros
    ----------
    xyz : array-like, shape (3,)
       Pontos no espaço 3D. Note que [x, y, z] estão empacotados numa variável só,
       para encurtar o programa.
    s, r, b : float s=10, r=28, b=8/3
       Parâmetros que definem o atrator de Lorenz.

    Retorna
    -------
    xyz_dot : array, shape (3,)
       Valores de dx/dt, dy/dt, dz/dt, nos pontos (x,y,z). Também retornam numa variável só.

    Do paper do Lorenz, sobre convecção:
    x (X) é proporcional à intesidade do movimento convectivo;
    y (Y) é proporcional à diferença de temperatura enre a corrente ascendente e a descendente;
    z (Z) é proporcional à diferença entre o perfil T(z) observado e uma reta.
    s = v/k  é o número de Prandtl (sigma) onde v (nu) é a viscosidade cinemática e k (kappa) é a 
        condutividade térmica.
    r = Ra/Rc, número de Rayleigh (Ra = (g \alpha H³ \Delta T) / (k v) ) dividido pelo 
        número de Rayleigh Crítico (Rc = \pi⁴)³/a². \alpha é o coef. de expansão térmica.
    b = 4/(1+a²) Constante. Mais detalhes no paper.

    """
    x, y, z = xyz
    x_dot = s*(y - x)
    y_dot = r*x - y - x*z
    z_dot = x*y - b*z
    return array([x_dot, y_dot, z_dot])

dt = 0.0025
num_steps = 25000
xyzs = empty((num_steps + 1, 3))  # Um espaço a mais para o valor inicial
b, r, eps = (8/3, 28, 0.01) # Se eps=0 o modelo é estável, pois é inicializado no centro do atrator
xyzs[0] = (-sqrt(b*(r-1+eps)), -sqrt(b*(r-1+eps)), r-1+eps) 
# xyzs[0] = (9., 9., 28.) # é igual ao do paper do Lorenz

# Avança no tempo, de dt em dt, calculando as derivadas no ponto corrente
# (i) para estimar o próximos valores, no tempo i+1. Viu como fica curtinho?!
for i in range(num_steps):
    xyzs[i + 1] = xyzs[i] + lorenz(xyzs[i]) * dt

# Terminado o loop, só falta plotar

# t1 é o tempo (arbitrário, para fazer figuras limpinhas)
t1 = arange(num_steps+1)/1000
cmap = plt.cm.gist_rainbow

fig = plt.figure(num=1, figsize=(6,10), clear=True)
fig.suptitle("Lorenz Attractor")
ax = fig.add_subplot(2, 1, 1, projection='3d')


# o * é para desempacotar x, y e z
line = ax.scatter(*xyzs.T, s=0.1, c=t1, cmap=cmap)
ax.set_xlabel("$w$")
ax.set_ylabel("$\Delta T$")
ax.set_zlabel("$T(z)-\overline{T}$")
#plt.colorbar(line, ax=ax, label='time')

bx = fig.add_subplot(6, 1, 4)
series = bx.scatter(t1, *xyzs[:,[0]].T, s=0.1, c=t1, cmap=cmap)
bx.set_ylabel("$w$")


cx = fig.add_subplot(6, 1, 5)
series = cx.scatter(t1, *xyzs[:,[1]].T, s=0.1, c=t1, cmap=cmap)
cx.set_ylabel("$\Delta T$")

dx = fig.add_subplot(6, 1, 6)
series = dx.scatter(t1, *xyzs[:,[2]].T, s=0.1, c=t1, cmap=cmap)
dx.set_xlabel("time")
dx.set_ylabel("$T(z)-\overline{T}$")

bx.set_xlim(0,t1.max())
cx.set_xlim(0,t1.max())
dx.set_xlim(0,t1.max())

bx.xaxis.set_major_locator(ticker.LinearLocator(numticks=6))
cx.xaxis.set_major_locator(ticker.LinearLocator(numticks=6))
dx.xaxis.set_major_locator(ticker.LinearLocator(numticks=6))

bx.xaxis.set_major_formatter(ticker.NullFormatter())
cx.xaxis.set_major_formatter(ticker.NullFormatter())

bx.grid('True')
cx.grid('True')
dx.grid('True')

plt.tight_layout()

plt.show()