Making Animations Quickly with Matplotlib Blitting

April 07, 2018 #data-visualization

A color image of the M101 Pinwheel Galaxy, a spiral galaxy, composed of multiple exposures taken by NASA and the European Space Agency.

Animations are a great way to show the passage of time in a plot. I have used animation to show how long my Raspberry Pis take to reboot and how the popularity of names changed in the US.

But making animations in matplotlib can take a long time. Not just to write the code, but waiting for it to run! The easiest, but slowest, way to make an animation is to redraw the entire plot every frame. Using this method, it took roughly 20 minutes to render a single animation for my names post! Fortunately there is a significantly faster alternative: matplotlib’s animation blitting. Blitting increased rendering speed by a factor of 20!

The Data

We will plot the spectrum of Supernova 2011fe from Pereira et al.¹ by the Nearby Supernova Factory.² The spectrum of a supernova tells us about what is going on in the explosion, so looking at a time series tells us how the explosion is evolving.

The data is available here. The notebook with all the code is here (rendered on Github). The code in the notebooks is complete, including doc strings and comments, while I have stripped down the examples below for clarity.

This is the animation we will be making:

It shows the amount of light (flux) the telescope saw as a function of the wavelength of light. The data was only sampled once every few days, so to make the animation smooth we will linearly interpolate the data. This is implemented by the function flux_from_day(day), which returns a numpy array of flux values for a specific day. The details of how the function works can be found in the notebook.

Blitting

Blitting breaks the animation into two components: the unchanging background elements, and the artist objects that are updated each frame. It requires us to write three functions:

init_fig(): draws the static background
frame_iter(): yields the frame_data needed to draw each update
update_artists(frame_data): takes frame_data and updates the artists

The artists that are updated each frame must be kept in an iterable container. A normal list will work, but a more convenient way to do this is using a namedtuple (which I discuss in detail in another post). This will let us access the different artists by name, for example artists.flux_line, instead of having to remember their index number.

init_fig Function

The init_fig() function draws the background of the animation. It takes no arguments and must return an iterable of the artists to be updated every frame, which in our case are contained in the namedtuple discussed above.

Our example function sets the labels, the title, and the range of the plot. It is here where we would draw anything else that is unchanging, like the legend, or some text labels, if we needed to. Here it is:

def init_fig(fig, ax, artists):
    """Initialize the figure, used to draw the first
    frame for the animation.
    """
    # Set the axis and plot titles
    ax.set_title("Supernova 2011fe Spectrum", fontsize=22)
    ax.set_xlabel("Wavelength [Å]", fontsize=20)
    FLUX_LABEL = "Flux [erg s$^{-1}$ cm$^{-2}$ Å$^{-1}$]"
    ax.set_ylabel(FLUX_LABEL, fontsize=20)

    # Set the axis range
    plt.xlim(3000, 10000)
    plt.ylim(0, 1.25e-12)

    # Must return the list of artists, but we use a pass
    # through so that they aren't created multiple times
    return artists

You will notice that I said the function takes no arguments, but I gave it three anyway. It’s hard to have no inputs (without using globals), but one trick is to use partial application, which I will demonstrate when we put it all together. The function must return the list of artists to update, but I find it’s easier to declare those outside of the function and then pass them in as an argument.

frame_iter Function

The frame_iter() function is a generator that returns the data needed to update the artist for each frame. It yields frame_data, which can be any sort of Python data type or object. This function also must take no arguments, and so like init_fig() we will use the partial application trick to bind the arguments.

Our function loops over the days relative to maximum light and returns the flux values from that day, as well as string of the day to update the text label.

def frame_iter(from_day, until_day):
    """Iterate through the days of the spectra and return
    flux and day number.
    """
    for day in range(from_day, until_day):
        flux = flux_from_day(day)
        # Yield events so the function can be looped over
        yield (flux, "Day: {day}".format(day))

update_artists Function

Once we have frame_iter() to generate the data for each frame, update_artists() is really simple. All it has to do is:

Unpack the frames_data.
Update the plot line and the text.

For the plot line we call .set_data() to insert the new values; for the text we call .set_text(). Our function is short:

def update_artists(frames, artists, lambdas):
    """Update artists with data from each frame."""
    flux, day = frames

    artists.flux_line.set_data(lambdas, flux)
    artists.day.set_text(day)

Lines and text are easy to update, but other plot objects (like histograms) are associated with multiple artists, which makes it harder to update them. Unfortunately, the only solution is to write a much more complicated update function for each type.

Putting it all together

Once we’ve written the three functions, it is pretty simple to make our animation:

Create the figure (fig) and axes (ax).
Create the list of artists, in this case a line (plt.plot) and some text (ax.text).
Partially apply the functions by binding inputs to them with partial.
Create the animation object (animation.FuncAnimation).
Save the animation as an .mp4 (anim.save).

Here are those steps in code:

# 1. Create the plot
fig, ax = plt.subplots(figsize=(12, 7))

# 2. Initialize the artists with empty data
Artists = namedtuple("Artists", ("flux_line", "day"))
artists = Artists(
    plt.plot([], [], animated=True)[0],
    ax.text(x=0.987, y=0.955, s=""),
)

# 3. Apply the three plotting functions written above
init = partial(init_fig, fig=fig, ax=ax, artists=artists)
step = partial(frame_iter, from_day=-15, until_day=25)
update = partial(update_artists, artists=artists,
                 lambdas=np.arange(3298, 9700, 2.5))

# 4. Generate the animation
anim = animation.FuncAnimation(
    fig=fig,
    func=update,
    frames=step,
    init_func=init,
    save_count=len(list(step())),
    repeat_delay=5000,
)

# 5. Save the animation
anim.save(
  filename='/tmp/sn2011fe_spectral_time_series.mp4',
  fps=24,
  extra_args=['-vcodec', 'libx264'],
  dpi=300,
)

The only tricky thing is the use of partial applications. Partial application binds some (or all) of the arguments to the function and creates a new function that takes fewer arguments. Essentially, it’s like setting a default value for the arguments.

For the update() function above, we use partial application to fix some of the arguments, while leaving the frame argument as one that still must be supplied at call time. To create the init() and step() functions above, we fully apply the parent functions, allowing the new functions to be called without any inputs.

A Little Extra

Of course, you can add a bit more to the plot, like the photometrics filters used:

But that would have made the example even harder to follow. If you’re interested, the notebook to generate that plot is here (rendered on Github).

Pereira et al., Spectrophotometric time series of SN 2011fe from the Nearby Supernova Factory, A&A 554, A27 (2013), doi: 10.1051/0004-6361/201221008 ↩
Aldering et al., Overview of the Nearby Supernova Factory, Proceedings Volume 4836, Survey and Other Telescope Technologies and Discoveries; (2002); doi: 10.1117/12.458107 ↩