Flow3D.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr 16 07:38:57 2018

@author: David Nolte
Introduction to Modern Dynamics, 2nd edition (Oxford University Press, 2019)

3D Flow examples: Lorenz, Rossler, Chua
"""
import numpy as np
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from scipy import integrate
from matplotlib import pyplot as plt

plt.close('all')
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1], projection='3d')
ax.axis('on')

# model_case 1 = Lorenz
# model_case 2 = Rossler
# model_case 3 = Chua
model_case = int(input('Enter Model Case (1-3)'));

def solve_lorenz(param, max_time=8.0, angle=0.0):

    if model_case == 1:
# Lorenz 3D flow 
        def flow_deriv(x_y_z, t0, sigma, beta, rho):
        #"""Compute the time-derivative of a Lorenz system."""
            x, y, z = x_y_z
            return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]
        model_title = 'Lorenz Attractor'

    elif model_case == 2:
# Rossler 3D flow 
        def flow_deriv(x_y_z, t0, sigma, beta, rho):
        #"""Compute the time-derivative of a Medio system."""
            x, y, z = x_y_z
            return [-y-z, x + sigma*y, beta + z*(x - rho)]
        model_title = 'Rossler Attractor'

    else:
# Chua 3D flow 
        def flow_deriv(x_y_z, t0, alpha, beta, c, d):
        #"""Compute the time-derivative of a Medio system."""
            x, y, z = x_y_z
            f = c*x + 0.5*(d-c)*(abs(x+1)-abs(x-1))
            return [alpha*(y-x-f), x-y+z, -beta*y]
        model_title = 'Chua Attractor'



    N=12
    colors = plt.cm.prism(np.linspace(0, 1, N))

 
    # Choose random starting points, uniformly distributed from -15 to 15
    np.random.seed(1)
    x0 = init1 + init2*np.random.random((N, 3))

    # Settle-down Solve for the trajectories
    t = np.linspace(0, max_time/4, int(250*max_time/4))
    x_t = np.asarray([integrate.odeint(flow_deriv, x0i, t, param)
                      for x0i in x0])
    
    # Solve for trajectories
    x0 = x_t[0:N,int(250*max_time/4)-1,0:3]
    t = np.linspace(0, max_time, int(250*max_time))
    x_t = np.asarray([integrate.odeint(flow_deriv, x0i, t, param)
                      for x0i in x0])
     
    
    # choose a different color for each trajectory
    # colors = plt.cm.viridis(np.linspace(0, 1, N))
    # colors = plt.cm.rainbow(np.linspace(0, 1, N))
    # colors = plt.cm.spectral(np.linspace(0, 1, N))
    colors = plt.cm.prism(np.linspace(0, 1, N))

    for i in range(N):
        x, y, z = x_t[i,:,:].T
        lines = ax.plot(x, y, z, '-', c=colors[i])
        plt.setp(lines, linewidth=0.5)

    ax.view_init(30, angle)
    plt.show()
    plt.title(model_title)
    plt.savefig('Flow3D')

    return t, x_t

if model_case == 1:
    param = (10, 8/3, 28)    # Lorenz
    ax.set_xlim((-25, 25))
    ax.set_ylim((-35, 35))
    ax.set_zlim((5, 55))
    max_time = 50.0
    init1 = -15
    init2 = 30
elif model_case == 2:
    param = (0.2, 0.2, 5.7)         # Rossler
    ax.set_xlim((-15, 15))
    ax.set_ylim((-15, 15))
    ax.set_zlim((0, 20))
    max_time = 200
    init1 = -15
    init2 = 30
else:
    param = (15.6, 28.0, -0.7, -1.14 )         # Chua
    ax.set_xlim((-3, 3))
    ax.set_ylim((-1, 1))
    ax.set_zlim((-3, 3))
    max_time = 100
    init1 = 0
    init2 = 0.1


t, x_t = solve_lorenz(param, max_time,angle=30)

plt.figure(2)
lines = plt.plot(t,x_t[1,:,0],t,x_t[1,:,1],t,x_t[1,:,2])
plt.setp(lines, linewidth=1)

for i in range(4):
    plt.figure(3)
    lines = plt.plot(x_t[i,:,0],x_t[i,:,1])
    plt.setp(lines,linewidth=0.5)

    plt.figure(4)
    lines = plt.plot(x_t[i,:,1],x_t[i,:,2])
    plt.setp(lines,linewidth=0.5)
    
    plt.figure(5)
    lines = plt.plot(x_t[i,:,0],x_t[i,:,2])
    plt.setp(lines,linewidth=0.5)