Animations in Matplotlib

In [ ]:
%pylab

We showed previously how to create basic animations in matplotlib using a simple canvas update:

In [ ]:
fig, ax = plt.subplots()
x = np.linspace(0, 10, 1000)
lines = ax.plot(x, np.sin(x), lw=2)
In [ ]:
for i in range(100):
    lines[0].set_ydata(np.sin(x + 0.1 * i))
    fig.canvas.draw()

The problem with this approach is that it can operate unpredictably between different backends, and may not even run on some! Also, if you'd like to save the animation to file, you must do a lot of things by hand. Thankfully, the Matplotlib team put together a very nice animation package in the version 1.1 release. Creating animations is fairly simple with this tool, as we'll see now.

Note that this framework is not without its own issues: in particular, animations will not work within the MacOSX backend (see, for example, https://github.com/matplotlib/matplotlib/issues/531)

Creating a Basic Animation

Creating a basic animation is a matter of initializing the plot, creating functions to update the frames, and passing these functions to an animation object.

We start by setting up a normal matplotlib figure:

In [ ]:
import matplotlib.pyplot as plt
import numpy as np

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))
line, = ax.plot([], [], lw=2)

Next we'll create an init() function and an animate() function.

The purpose of the init() function is to set the background of the animation: it should essentially hide any plot elements that you don't want to be shown in every frame.

The purpose of the animate() function is to update the plot elements for each frame.

In [ ]:
# initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,
In [ ]:
# animation function.  This is called sequentially
x = np.linspace(0, 2, 1000)

def animate(i):
    y = np.sin(2 * np.pi * (x - 0.01 * i))
    line.set_data(x, y)
    return line,

Creating the animation now is a matter of passing these initialization and frame-step functions to the animator:

In [ ]:
from matplotlib import animation

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=200, interval=20, blit=True)

Saving the Animation

To save the animation, you'll need one of several additional programs available on your system. The animation library works with both `ffmpeg,mencoder, andavconv, and the current master of matplotlib also works withimagemagick`` to natively create animated gifs.

For more information on these, see http://matplotlib.sourceforge.net/api/animation_api.html

The save command looks something like this:

In [ ]:
anim.save('basic_animation.mp4', fps=30)

3D Animation: The Lorenz Equations

The Lorenz equations are an interesting set of nonlinear differential equations which pop up in a surprising number of physical systems. Here we'll use scipy to solve the equations, and matplotlib to animate the results in 3D. The full script can be found in examples/lorenz_animation.py

The Lorenz Equations are as follows:

$dx/dt = \sigma(y−x)$

$dy/dt = x(\rho−z)−y$

$dz/dt = x y − \beta z$

with free parameters $\sigma$, $\rho$, and $\beta$. We can use scipy to solve these equations relatively easily:

In [ ]:
import numpy as np
from scipy import integrate

# define the lorenz derivatives
def lorenz_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]

# Choose random starting points, uniformly distributed from -15 to 15
N_trajectories = 20
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(lorenz_deriv, x0i, t)
                  for x0i in x0])

We can plot the resulting trajectories in 3D using the mplot3d toolkit:

In [ ]:
from mpl_toolkits.mplot3d import Axes3D

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

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

for i in range(N_trajectories):
    x, y, z = x_t[i].T
    ax.plot(x, y, z, c=colors[i])

Animating the Result

We can use this machinery the same way we have previously to animate these trajectories in 3D. You should recognize most of this from what we've seen earlier:

In [ ]:
from mpl_toolkits.mplot3d import Axes3D

# 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 = []
pts = []
for c in colors:
    lines += ax.plot([], [], [], '-', c=c)
    pts += ax.plot([], [], [], 'o', c=c)

# 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:])

    # rotate the point of view
    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)

The following command can be run to save this as an mp4:

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

Other Fun Examples

I've included several other fun animation examples with this tutorial, which you can run from the examples directory and inspect via the %load command:

Double Penulum

This is an example I adapted from the matplotlib documentation, and wrote about here: http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/

In [ ]:
%run ../examples/double_pendulum.py

Particles in a Box

This is another example I created and showed here: http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/

In [ ]:
%run ../examples/particle_box.py

Schrodinger Equation

This example is slightly more involved, and uses a half-step Schrodinger Equation algorithm to solve for the evolution of a wave function within a given potential. Details on the algorithm can be seen here: http://jakevdp.github.com/blog/2012/09/05/quantum-python/

In [ ]:
%run ../examples/animate_schrodinger.py