``````

In [1]:

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('bmh')
from numba import jit
from scipy.integrate import odeint
from time import time

``````
``````

In [2]:

# Just-in-time compiling these two functions
# seems to gain about a factor 2 in speed.
# (Note that the first time the function is called,
# it is slower, since compilation takes place.)
# (Note also that you should never define and use
# a just-in-time compiled function in the same cell
# in a notebook, as it will be compiled every time)

@jit
def lorenz(x, sigma, rho, beta):
# The Lorenz system (see Ch. 9 in Strogatz)
xdot = sigma*(x[1] - x[0])
ydot = x[0]*(rho - x[2]) - x[1]
zdot = x[0]*x[1] - beta*x[2]
return np.array([xdot, ydot, zdot])

@jit
def f(x, t):
# Wrapper function for the lorenz equations
sigma = 10.0
rho   = 28.0
beta  = 8/3
return lorenz(x, sigma, rho, beta)

``````
``````

In [3]:

x0 = [  5.62590715,  21.47606796,  26.79759756]
T  = np.linspace(0, 20, 10000)

tol = 1e-12
tic = time()
X   = odeint(f, x0, T, atol = tol, rtol = tol)
toc = time()

print('Calculating trajectory took %s seconds' % (toc - tic))

``````
``````

Calculating trajectory took 0.33465099334716797 seconds

``````
``````

In [4]:

# Reproducing approximately figure 9.4.1 from Strogatz
fig = plt.figure(figsize = (12,8))
plt.plot(X[:,1], X[:,2], lw = 1)
plt.xlabel('y')
plt.ylabel('z')

``````
``````

Out[4]:

<matplotlib.text.Text at 0x10ce7d2b0>

``````
``````

In [5]:

# Reproducing approximately figure 9.4.2 from Strogatz
fig = plt.figure(figsize = (12,6))
plt.plot(T, X[:,2])
plt.xlabel('t')
plt.ylabel('z')

``````
``````

Out[5]:

<matplotlib.text.Text at 0x10fb07470>

``````
``````

In [6]:

# A function to find a list of local maxima
@jit
def maxima(x):
mask = np.zeros(x.size, dtype = np.bool)
for i in range(1, x.size - 1):
if (x[i] > x[i-1]) and (x[i] > x[i+1]):

``````
``````

In [7]:

# Get a longer series of data
x0 = [  5.62590715,  21.47606796,  26.79759756]
T  = np.linspace(0, 1000, 1000000)
tol = 1e-12
X   = odeint(f, x0, T, atol = tol, rtol = tol)
# Get list of local maxima
m = maxima(X[:,2])

``````
``````

In [8]:

# Reproducing approximately figure 9.4.3 from Strogatz
# Scatter plot z_n+1 as a function of z_n
# (Remove first and last points)
N = m.size
fig = plt.figure(figsize = (12,8))
plt.scatter(m[1:N-2], m[2:N-1], marker = '.', lw = 0)
plt.plot([30,48], [30,48], lw = 1, c = 'k', alpha = 0.5)
plt.xlim(30,48)
plt.ylim(30,48)

``````
``````

Out[8]:

(30, 48)

``````
``````

In [ ]:

``````