Matplotlib animations the easy way

Creating animations with Python’s Matplotlib is quick and easy once you know how to do it. However, when learning I found the tutorials and examples online either daunting, overly sophisticated, or lacking explanation. In many cases all I need is a quick-and-dirty script that works, rather than longer code that adheres to best practices.

Why animate?

Exploring datasets is a big part of what many scientists do these days. In many cases these datasets will have more than two dimensions; for example, temperature or salinity in an ocean circulation model has four dimensions: x, y, z, t. It’s futile to try and display these in a single plot. That’s where animation can help.

An animated line in six steps

This example walks through how to create the animation below in six steps. However, the first four steps will involve nothing new to anyone who has made a plot using Matplotlib.

Each step contains a few lines of code that you can copy and paste, but a script with all the code for all examples can be found here.

animation_example
This animation requires less than 20 lines of code
Step one: import the necessary modules
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
# plt.style.use('ggplot')

The first two lines will be familiar to anyone who has used Python for science, and the third line is obviously specific to animation. I’ve left the final line commented as it isn’t necessary and will not work if your matplotlib version is <1.5.

Step two: set up the plotting area
fig, ax = plt.subplots(figsize=(5, 3))
ax.set(xlim=(-3, 3), ylim=(-1, 1))

The first line sets up the figure and its axis, and the second line fixes the axis limits. Setting the limits in advance stops any rescaling of the limits that may make the animation jumpy and unusable.

Step three: create some data to plot
x = np.linspace(-3, 3, 91)
t = np.linspace(1, 25, 30)
X2, T2 = np.meshgrid(x, t)

sinT2 = np.sin(2*np.pi*T2/T2.max())
F = 0.9*sinT2*np.sinc(X2*(1 + sinT2))

F is a 2D array comprising some arbitrary data to be animated. Well, it’s not actually arbitrary, it’s set up to produce a perfectly looping gif.

Step four: plot the first line
line = ax.plot(x, F[0, :], color='k', lw=2)[0]

This sets up a line object with the desired attributes, which in this case are that it’s coloured black and has a line weight of 2. Note the [0] at the end. This is necessary because the plot command returns a list of line objects. Here, we are only plotting a single line, so we simply want the first (i.e., zeroth) object in the list of lines.

Be sure to use ax.plot(…), not plt.plot(…)

Step five: create a function to update the line
def animate(i):
    line.set_ydata(F[i, :])

This function needs one argument. The only command in the function is to change the object’s data. See later for other things that can be changed.

Step six: call FuncAnimation and show
anim = FuncAnimation(
    fig, animate, interval=100, frames=len(t)-1)

plt.draw()
plt.show()

FuncAnimation requires two arguments: the figure object ‘fig’ and animation function ‘animate’. The interval argument is optional and sets the interval between frames in milliseconds. The frames argument is needed only if the animation is to be exported. The final two lines may or may not be necessary depending on whethere you’re working interactively (e.g. with IPython).

Be sure to assign FuncAnimation to a variable (here it’s anim).

Optional step: export the animation

Exporting the animation should be as simple as

anim.save('filename.mp4')

or, say

anim.save('filename.gif', writer='imagemagick')

However, depending on your computer and filetype, some configuration may be required to ensure the export works correctly.

Beyond a line plot

Animating any other type of plot is as simple as the example above. Typically, the only differences involve getting the equivalent of the line object, and  changing the animate function. Here are examples of steps four and five for the most common types of plots.

We’ll be needing a three dimensional array G for some of these examples.

x = np.linspace(-3, 3, 91)
t = np.linspace(0, 25, 30)
y = np.linspace(-3, 3, 91)
X3, Y3, T3 = np.meshgrid(x, y, t)
sinT3 = np.sin(2*np.pi*T3 /
               T3.max(axis=2)[..., np.newaxis])
G = (X3**2 + Y3**2)*sinT3

Pcolor

cax = ax.pcolormesh(x, y, G[:-1, :-1, 0],
                    vmin=-1, vmax=1, cmap='Blues')
fig.colorbar(cax)

def animate(i):
     cax.set_array(G[:-1, :-1, i].flatten())

Note that the length of G in each dimension must be one shorter than x and y, and that a flattened array must be passed to set_array.

pcolor

Scatter

ax.set(xlim=(-3, 3), ylim=(-1, 1))
scat = ax.scatter(x[::3], F[0, ::3])

def animate(i):
    y_i = F[i, ::3]
    scat.set_offsets(np.c_[x[::3], y_i])

Note the set_offsets must be passed an N×2 array. Here we use np.c_[] for this purpose. The use of [::3] reduces the density of scatter points.

scatter

Quiver

ax.set(xlim=(-4, 4), ylim=(-4, 4))

# Plot every 20th arrow
step = 20
x_q, y_q = x[::step], y[::step]

# Create U and V vectors to plot
U = G[::step, ::step, :]
V = np.roll(U, shift=3, axis=2)

qax = ax.quiver(x_q, y_q,
                U[..., 0], V[..., 0],
                scale=100)

def animate(i):
    qax.set_UVC(U[..., i], V[..., i])

quiver

Contour

contour_opts = {'levels': np.linspace(-9, 9, 10),
                'cmap':'RdBu', 'lw': 2}
cax = ax.contour(x, y, G[..., 0], **contour_opts)

def animate(i):
    ax.collections = []
    ax.contour(x, y, G[..., i], **contour_opts)

Contour plots are a little different. We don’t use a set_… method, but instead simply redraw the contour plot for each frame. The first line of the function gets rid of existing contours before plotting new ones.

Using the contour_opts dict is a handy trick to avoid specifying the same keyword arguments in both calls to contour.

contour

Changing labels and text

ax.set(xlim=(-1, 1), ylim=(-1, 1))
string_to_type = 'abcdefghijklmnopqrstuvwxyz0123'
label = ax.text(0, 0, string_to_type[0],
                ha='center', va='center',
                fontsize=12)

def animate(i):
    label.set_text(string_to_type[:i+1])
    ax.set_ylabel('Time (s): ' + str(i/10))
    ax.set_title('Frame ' + str(i))

labels.gif

Changing other properties

With the exception of contours, every animate function requires finding the appropriate set_… method. The following list comprehension will give you all the potential set methods for, say, a line object. (You have to create this first.)

[x for x in dir(line) if 'set_' in x]

This often gives you a long list of potential values to set, but usually a quick scan through the list is all that is needed to find the method you want.

Advertisements

Author: hugke729

PhD student in physical oceanography

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s