Exercise 9.2: Numerically solving differential equations

---

import numpy as np

import bokeh.io
import bokeh.plotting

bokeh.io.output_notebook()

Loading BokehJS ...

For this exercise, we will use Euler’s method to simulate a classic set of differential equations that can exhibit chaotic solutions known as the Lorenz attractor. As a warm-up example, we will simulate bacterial growth.

Bacterial growth can be modeled by the differential equation

\begin{align} \frac{\mathrm{d}n}{\mathrm{d}t} = k n, \end{align}

where \(n\) is the number of bacteria and \(k\) is the growth rate. The idea here is that the number of bacteria will grow faster the more bacteria we have, because there are more to divide. Analytically, we know the solution to this differential equation is

\begin{align} n(t) = n_0 \mathrm{e}^{kt}, \end{align}

i.e., exponential growth. But suppose we did not know how to derive that. We could simulate the differential equation. We do this by discretizing time. Instead of a derivative, we have a change in \(n\) over a change in time \(t\).

\begin{align} \frac{\mathrm{d}n}{\mathrm{d}t} \approx \frac{\Delta n}{\Delta t} = k n. \end{align}

Let’s say we know \(n\) and time zero, \(n(0)\). Then \(n\) at time \(t = \Delta t\) is

\begin{align} n(\Delta t) \approx n(0) + \Delta n = n(0) + \Delta t\, k n(0). \end{align}

More generally, we can write

\begin{align} \frac{\mathrm{d}n}{\mathrm{d}t} = f(n), \end{align}

and

\begin{align} n(t+\Delta t) \approx n(t) + \Delta t\,f(n). \end{align}

So, we can instruct Python to take our current value of \(n\), and then add \(\Delta t\) times \(f(n)\) to get our new \(n\) at a time just a bit later on, at \(t + \Delta t\). We repeat this over and over again to move forward in time. Let’s code that up!

# Specify parameter
k = 1

# Specify my little time step
delta_t = 0.01

# Make an array of time points, evenly spaced up to 10
t = np.arange(0, 10, delta_t)

# Make an array to store the number of bacteria
n = np.empty_like(t)

# Set the initial number of bacteria
n[0] = 100

# Write a for loop to keep updating n as time goes on
for i in range(1, len(t)):
    n[i] = n[i - 1] + delta_t * k * n[i - 1]

Ok! We just computed the time points and the number of bacteria, so we can just plot the result!

p = bokeh.plotting.figure(
    height=300,
    width=450,
    x_axis_label="time (units of 1/k)",
    y_axis_label="number of bacteria",
)

p.line(x=t, y=n, line_width=2)

bokeh.io.show(p)

And there is the famous exponential growth!

This time stepping method is called Euler’s method, and what we’re doing is called numerical solution of a differential equation.

a) Now it’s time to solve the Lorenz attractor. In this case, there are three differential equations.

\begin{align} &\frac{\mathrm{d}x}{\mathrm{d}t} = \sigma(y - x),\\[1em] &\frac{\mathrm{d}y}{\mathrm{d}t} = x(\rho - z) - y,\\[1em] &\frac{\mathrm{d}z}{\mathrm{d}t} = xy - \beta z. \end{align}

Your task in this exercise is to numerically solve these two differential equations together and then plot the result. Use the following parameter values, which were the ones originally used by Lorenz.

sigma = 10
beta = 8/3
rho = 28

Start at time \(t = 0\) with \(x = y = z= 1\).

Even though there are now three differential equations, the procedure is the same, you update each by adding \(\Delta t\) times the respective derivative. You should integrate these equations until time \(t = 60\). I would recommend a step size of 0.001.

b) Plot your solution showing \(x\) versus time.

c) Solve again, this time with \(x(0) = 1.0001\), just 0.01% different from your previous starting point. Plot this solution on the same plot as your did in part (b) so you can compare the trajectories for just slightly different initial conditions.

d) Now make three plots. Plot \(x\) vs \(y\), \(x\) vs \(x\), and \(y\) vs \(z\). Look at the result. Cool, right?

e) [Bonus] Euler’s method is probably the simplest way to solve differential equations, and is by no means the best. SciPy has an ODE solver, scipy.integrate.odeint() that uses the more sophisticated and robust methods for solving systems of ODEs. Read the documentation about how scipy.integrate.odeint() works and use it to solve the Lorenz attractor system of ODEs.

This last part is tough; I’m not giving you directions, and you are kind of on your own to read the documentation and figure it out.

Solution

---

import scipy.integrate

a) We take the same approach as in the exponential growth example. We just have to update three variables at each time step, x, y, and z.

# Specify parameters
sigma = 10
beta = 8 / 3
rho = 28
delta_t = 0.001
t_end = 60
x_0 = y_0 = z_0 = 1


def solve_lorenz_euler(sigma, beta, rho, x_0, y_0, z_0, delta_t, t_end):
    # Time points
    t = np.arange(0, t_end, delta_t)

    # Make arrays to store rabbit and fox populations
    x = np.empty_like(t)
    y = np.empty_like(t)
    z = np.empty_like(t)

    # Set initial conditions
    x[0] = x_0
    y[0] = y_0
    z[0] = z_0

    # Write a for loop to keep updating r and f as time goes on
    for i in range(1, len(t)):
        x[i] = x[i - 1] + delta_t * sigma * (y[i - 1] - x[i - 1])
        y[i] = y[i - 1] + delta_t * (x[i - 1] * (rho - z[i - 1]) - y[i - 1])
        z[i] = z[i - 1] + delta_t * (x[i - 1] * y[i - 1] - beta * z[i - 1])

    return t, x, y, z


t, x, y, z = solve_lorenz_euler(
    sigma, beta, rho, x_0, y_0, z_0, delta_t, t_end
)

b) Ok, let’s see what we got! We will only plot every tenth data point so as not to overwhelm the browser.

p = bokeh.plotting.figure(
    height=200,
    width=600,
    x_axis_label='t',
    y_axis_label='x, y, z'
)

p.line(
    x=t[::10],
    y=x[::10],
    line_width=1,
)


bokeh.io.show(p)

c) Now, let’s change x_0, solve again, and add to the plot.

x_0 = 1.0001

t, x, y, z = solve_lorenz_euler(sigma, beta, rho, x_0, y_0, z_0, delta_t, t_end)

p.line(
    x=t[::10],
    y=x[::10],
    line_width=1,
    color='orange'
)

bokeh.io.show(p)

After progressing together for some time, the trajectories become completely uncorrelated after a while. This is an example of chaos; very similar, but slightly different, initial conditions give completely different dynamics.

d) Now, let’s make plots in the appropriate planes. We’ll write a function to do it, since we’ll use it again. We’ll allow for a stride kwarg to choose what fraction of the data points we want to use in the plots.

def plot_lorenz(x, y, z, stride=1):
    p_xy = bokeh.plotting.figure(
        height=200, width=200, x_axis_label="x", y_axis_label="y"
    )

    p_xz = bokeh.plotting.figure(
        height=200,
        width=200,
        x_axis_label="x",
        y_axis_label="z",
        x_range=p_xy.x_range,
    )

    p_yz = bokeh.plotting.figure(
        height=200,
        width=200,
        x_axis_label="y",
        y_axis_label="z",
        x_range=p_xy.y_range,
        y_range=p_xz.y_range,
    )

    p_xy.line(x=x[::stride], y=y[::stride])
    p_xz.line(x=x[::stride], y=z[::stride])
    p_yz.line(x=y[::stride], y=z[::stride])

    return bokeh.layouts.gridplot([[p_xy, None], [p_xz, p_yz]])


bokeh.io.show(plot_lorenz(x, y, z, stride=10))

e) Now, let’s use scipy.integrate.odeint() to do the solution. We need to specify a function to compute the derivatives.

def dxyz_dt(xyz, t, sigma, beta, rho):
    """Right hand side of Lorenz attractor"""
    # Unpack x, y, and z
    x, y, z, = xyz

    # Compute derivatives
    dx_dt = sigma * (y - x)
    dy_dt = x * (rho - z) - y
    dz_dt = x * y - beta * z

    return np.array([dx_dt, dy_dt, dz_dt])

Now we have to specify the initial conditions.

xyz_0 = np.array([1, 1, 1])

And specify the time points we want.

t = np.linspace(0, 60, 5000)

And then we just stuff the system into scipy.integrate.odeint(). We have to be sure to pass the additional arguments that dy_dt() takes as a tuple.

y = scipy.integrate.odeint(dxyz_dt, xyz_0, t, args=(sigma, beta, rho))

# Unpack coordinates
x, y, z = y.transpose()

# Make a plot
bokeh.io.show(plot_lorenz(x, y, z))

Computing environment

%load_ext watermark
%watermark -v -p numpy,scipy,bokeh,jupyterlab

CPython 3.7.7
IPython 7.16.1

numpy 1.18.5
scipy 1.5.0
bokeh 2.1.1
jupyterlab 2.1.5