One of the things I really enjoy about Python is how easy it makes it to solve interesting problems and visualize those solutions in a compelling way. I've done several posts on creating animations using matplotlib's relatively new animation toolkit: (some examples are a chaotic double pendulum, the collisions of particles in a box, the time-evolution of a quantum-mechanical wavefunction, and even a scene from the classic video game, Super Mario Bros.).

Recently, a reader commented asking whether I might do a 3D animation example. Matplotlib has a decent 3D toolkit called mplot3D, and though I haven't previously seen it used in conjunction with the animation tools, there's nothing fundamental that prevents it.

At the commenter's suggestion, I decided to try this out with a simple example of a chaotic system: the Lorenz equations.

Solving the Lorenz System

The Lorenz Equations are a system of three coupled, first-order, nonlinear differential equations which describe the trajectory of a particle through time. The system was originally derived by Lorenz as a model of atmospheric convection, but the deceptive simplicity of the equations have made them an often-used example in fields beyond atmospheric physics.

The equations describe the evolution of the spatial variables $x$, $y$, and $z$, given the governing parameters $\sigma$, $\beta$, and $\rho$, through the specification of the time-derivatives of the spatial variables:

${\rm d}x/{\rm d}t = \sigma(y - x)$

${\rm d}y/{\rm d}t = x(\rho - z) - y$

${\rm d}z/{\rm d}t = xy - \beta z$

The resulting dynamics are entirely deterministic giving a starting point $(x_0, y_0, z_0)$ and a time interval $t$. Though it looks straightforward, for certain choices of the parameters $(\sigma, \rho, \beta)$, the trajectories become chaotic, and the resulting trajectories display some surprising properties.

Though no general analytic solution exists for this system, the solutions can be computed numerically. Python makes this sort of problem very easy to solve: one can simply use Scipy's interface to ODEPACK, an optimized Fortran package for solving ordinary differential equations. Here's how the problem can be set up:

import numpy as np
from scipy import integrate

# Note: t0 is required for the odeint function, though it's not used here.
def lorentz_deriv((x, y, z), t0, sigma=10., beta=8./3, rho=28.0):
    """Compute the time-derivative of a Lorenz system."""
    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]

x0 = [1, 1, 1]  # starting vector
t = np.linspace(0, 3, 1000)  # one thousand time steps
x_t = integrate.odeint(lorentz_deriv, x0, t)

That's all there is to it!

Visualizing the results

Now that we've computed these results, we can use matplotlib's animation and 3D plotting toolkits to visualize the trajectories of several particles. Because I've described the animation tools in-depth in a previous post, I will skip that discussion here and jump straight into the code:

import numpy as np
from scipy import integrate

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import cnames
from matplotlib import animation

N_trajectories = 20


def lorentz_deriv((x, y, z), t0, sigma=10., beta=8./3, rho=28.0):
    """Compute the time-derivative of a Lorentz system."""
    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]


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

# Solve for the trajectories
t = np.linspace(0, 4, 1000)
x_t = np.asarray([integrate.odeint(lorentz_deriv, x0i, t)
                  for x0i in x0])

# Set up figure & 3D axis for animation
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1], projection='3d')
ax.axis('off')

# choose a different color for each trajectory
colors = plt.cm.jet(np.linspace(0, 1, N_trajectories))

# set up lines and points
lines = sum([ax.plot([], [], [], '-', c=c)
             for c in colors], [])
pts = sum([ax.plot([], [], [], 'o', c=c)
           for c in colors], [])

# prepare the axes limits
ax.set_xlim((-25, 25))
ax.set_ylim((-35, 35))
ax.set_zlim((5, 55))

# set point-of-view: specified by (altitude degrees, azimuth degrees)
ax.view_init(30, 0)

# initialization function: plot the background of each frame
def init():
    for line, pt in zip(lines, pts):
        line.set_data([], [])
        line.set_3d_properties([])

        pt.set_data([], [])
        pt.set_3d_properties([])
    return lines + pts

# animation function.  This will be called sequentially with the frame number
def animate(i):
    # we'll step two time-steps per frame.  This leads to nice results.
    i = (2 * i) % x_t.shape[1]

    for line, pt, xi in zip(lines, pts, x_t):
        x, y, z = xi[:i].T
        line.set_data(x, y)
        line.set_3d_properties(z)

        pt.set_data(x[-1:], y[-1:])
        pt.set_3d_properties(z[-1:])

    ax.view_init(30, 0.3 * i)
    fig.canvas.draw()
    return lines + pts

# instantiate the animator.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=500, interval=30, blit=True)

# Save as mp4. This requires mplayer or ffmpeg to be installed
#anim.save('lorentz_attractor.mp4', fps=15, extra_args=['-vcodec', 'libx264'])

plt.show()

The resulting animation looks something like this:

Notice that there are two locations in the space that seem to draw-in all paths: these are the so-called "Lorenz attractors", and have some interesting properties which you can read about elsewhere. The qualitative characteristics of these Lorenz attractors vary in somewhat surprising ways as the parameters $(\sigma, \rho, \beta)$ are changed. If you are so inclined, you may wish to download the above code and play with these values to see what the results look like.

I hope that this brief exercise has shown you the power and flexibility of Python for understanding and visualizing a large array of problems, and perhaps given you the inspiration to explore similar problems.

Happy coding!