using Plots

The Lorenz Attractor

(see Swanson p. 452)

Solve the Lorenz equations:

$$ \begin{align} \dot{x} &= \sigma \, (y-x) \\ \dot{y} &= x(\rho - z) - y \\ \dot{z} &= xy - \beta z \end{align} $$

with $\sigma = 10$, $\beta = 8/3$. There is a critical value of $\rho$:

$$ \rho_c = 470/19 \approx 24.7 $$ for chaos. We can try $\rho = (10, 28)$ to see.

function df(xv; sigma=10.0, beta=8/3, rho=10)

    x, y, z = xv

    dx = sigma*(y-x)
    dy = x*(rho-z)-y
    dz = x*y - beta*z

    return [dx, dy, dz]
end

df (generic function with 1 method)

function df(xv; sigma=10.0, beta=8/3, rho=10)

    x, y, z = xv

    dx = sigma*(y-x)
    dy = x*(rho-z)-y
    dz = x*y - beta*z

    return [dx, dy, dz]
end

# dt = 0.04 Euler is not accurate; but RK4 works
dt = 0.02

Nmax = 5000

t = 0.0
xv = [1.0, -2.0, 1.0]

data = zeros(Nmax, 4)

data[1, :] = [t, xv...]'

for i1 in 2:Nmax

    t += dt

    # simple Euler
    # xv +=  dt * df(xv)

    # RK4 on the fly
    
    df1 = df(xv)
    df2 = df(xv+dt*0.5*df1)
    df3 = df(xv+dt*0.5*df2)    
    df4 = df(xv+dt*df3)
    
    xv += dt * (df1 + 2*df2 + 2*df3 + df4)/6

    data[i1, :] = [t, xv...]'

end

# Plot the xz projection
p1 = plot(data[:, 2], data[:, 4], 
     label="", xlabel="x", ylabel="z",
     # linewidth=0.8,
     alpha=0.1,
     marker=:circle,
     title="Lorenz Attractor rho=10 (projected to xz plane)")

# Create a custom colormap along the rainbow
mycolor = cgrad([:red, :orange, :yellow, :green, :blue, :indigo, :purple])

p1 = plot(data[:, 2], data[:, 4], 
          line_z = 1:size(data, 1), 
          colormap = mycolor,
          linewidth = 2,
          legend = false,
          alpha = 0.2,
          xlabel = "x", ylabel = "z",
          title = "Lorenz Attractor rho=10 (xz plane)")

scatter!((data[1,2], data[1,4]), color=:red, label="start")
scatter!((data[end,2], data[end,4]), color=:purple, label="finish")

function df(xv; sigma=10.0, beta=8/3, rho=28)

    x, y, z = xv

    dx = sigma*(y-x)
    dy = x*(rho-z)-y
    dz = x*y - beta*z

    return [dx, dy, dz]
end

dt = 0.025

Nmax = 6500
data = zeros(Nmax, 4)

t = 0.0
xv = [0, 1.0, -0.5]

data[1, :] = [t, xv...]'


for i1 in 2:Nmax

    t += dt

    # simple Euler
    # xv +=  dt * df(xv...)

    # RK4 on the fly
    df1 = df(xv)
    df2 = df(xv+dt*0.5*df1)
    df3 = df(xv+dt*0.5*df2)    
    df4 = df(xv+dt*df3)
    xv += dt * (df1 + 2*df2 + 2*df3 + df4)/6

    data[i1, :] = [t, xv...]'

end

# Plot the xz projection
p1 = plot(data[:, 2], data[:, 4], 
     label="", xlabel="x", ylabel="z",
     # linewidth=0.8,
     alpha=0.1,
     marker=:circle,
     title="Lorenz Attractor rho=10 (projected to xz plane)")

# Create a custom colormap along the rainbow
mycolor = cgrad([:red, :orange, :yellow, :green, :blue, :indigo, :purple])

p1 = plot(data[:, 2], data[:, 4], 
          line_z = data[:,1], 
          colormap = mycolor,
          linewidth = 2,
          legend = true,
          alpha = 0.2,
          xlabel = "x", ylabel = "z",
          title = "Lorenz Attractor rho=28 (xz plane)",
          label="")

scatter!((data[1,2], data[1,4]), color=:red, label="start")
scatter!((data[end,2], data[end,4]), color=:purple, label="finish")

function df(xv; sigma=10.0, beta=8/3, rho=28.0)

    x, y, z = xv

    dx = sigma*(y-x)
    dy = x*(rho-z)-y
    dz = x*y - beta*z

    return [dx, dy, dz]
end

dt = 0.025
Nmax = 120
data = zeros(Nmax, 3)

t = 0.0
xv = [0.1, 0.2, -0.5]

data[1, :] = [t, xv[1], xv[3]]'

itrack = 1
while true

    t += dt

    # simple Euler
    # xv +=  dt * df(xv)

    # RK4 on the fly
    df1 = df(xv)
    df2 = df(xv+dt*0.5*df1)
    df3 = df(xv+dt*0.5*df2)    
    df4 = df(xv+dt*df3)
    xv1 = xv + dt * (df1 + 2*df2 + 2*df3 + df4)/6

    if xv1[2] * xv[2] < 0

        # simple mid-point
        # xr = (xv[1] + xv1[1])/2.0
        # zr = (xv[3] + xv1[3])/2.0

        # with linear interp
        m = xv[2] / (xv[2] - xv1[2])
        xr = xv[1] + m * (xv1[1] - xv[1])
        zr = xv[3] + m * (xv1[3] - xv[3])

        itrack += 1
        data[itrack, :] = [t, xr, zr]'
        
        # println([t, xr, zr, xv[2]])
    else
    end

    if itrack >= Nmax
            break
    else
    end
    xv .= xv1

end


# Create a custom colormap along the rainbow
mycolor = cgrad([:red, :orange, :yellow, :green, :blue, :indigo, :purple])

p1 = scatter(data[:, 2], data[:, 3], 
          marker_z = 1:length(data[:, 1]), 
          colormap = mycolor,
          linewidth = 2,
          legend = false,
          alpha = 0.2,
          xlabel = "x", ylabel = "z",
          title = "Poincaré Section y=0")

scatter!((data[1,2], data[1,3]), color=:red, label="start")
scatter!((data[end,2], data[end,3]), color=:purple, label="finish")