lorentz.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

def func(r, t, params):
    x, y, z = r
    a, b, c = params
    
    # Differential Equations
    dx = a*(y-x)
    dy = ((c-a)*x) - (x*z) + (c*y)
    dz = (x*y)-(b*z)
    
    deriv = [dx, dy, dz]
    return deriv

# Initial Params a, b, c
params = [40, 5, 30]


# Initial Conditions x0, y0, z0
r0 = [1, 1, 1]
# Time array start, stop, increment
t = np.arange(0., 10, 0.001)
n = int(10/0.001)
s = 10 # Segment length

# Call ODE solver
psol = odeint(func, r0, t, args=(params, ))

sol0 = psol[:, 0]
sol1 = psol[:, 1]
sol2 = psol[:, 2]

fig1 = plt.figure()

T=np.linspace(0,1,np.size(sol0))**2
tlen = len(T)-1

ax4 = fig1.add_subplot(111)
for i in range(0,n-s,s):
    ax4.plot(sol0[i:i+s+1],sol1[i:i+s+1],color=(0,T[i],T[i]))
ax4.set_xlabel("x")
ax4.set_ylabel("y")

ax5 = fig1.add_subplot(111)
for i in range(0,n-s,s):
    ax4.plot(sol1[i:i+s+1],sol2[i:i+s+1],color=(T[i],T[i],0))
ax5.set_xlabel("y")
ax5.set_ylabel("z")

ax6 = fig1.add_subplot(111)

for i in range(0,n-s,s):
    ax4.plot(sol0[i:i+s+1],sol2[i:i+s+1],color=(T[i],0,T[i]))
ax6.set_xlabel("r")
ax6.set_ylabel("r")

plt.draw()
plt.show()