Elaborate Matplotlib animations

Matplotlib, Python’s primary scientific plotting library, provides tools to make many elaborate plots, graphs, and diagrams. Many of these can be animated, but the process isn’t always intuitive. The hardest part is learning how to animate a simple line plot (here’s my easy way). Beyond that, the steps to creating most animations tend to be similar.

The examples below demonstrate the particular methods needed to animate common types of plot. Here I focus on the key components needed for updating each frame of the animation. The full code for the examples is here. It includes liberal and arbitrary use of sines and cosines so as to produce looping gifs.

Scatter with variable colour, position, and size

Here we are animating 8 × 8 = 64 points. X3 and Y3, the locations of the circles, are each arrays of size (N_f, 8, 8) where N_f is the number of frames. S (the sizes of the circles) is the same size as X3, whereas C (the colours of the circles) has size (N_f, 8, 8, 4).

def animate(i):
    Xi, Yi = X3[i, :, :].flatten(), Y3[i, :, :].flatten()
    scat.set_offsets(np.c_[Xi, Yi])
    scat.set_sizes(S[i, :, :].flatten())
    scat.set_facecolors(C[i, :, :].reshape(-1, 4))

• set_offsets alters the location of the points. It accepts a 64 × 2 array of (x, y) coordinates. np.c_ is a handy tool to combine two 1D arrays appropriately.
• set_sizes does what it sounds like. It accepts a 1D array of length 64.
• set_facecolors also does what it sounds like. It accepts a 64 × 4 array, where each row is (R, G, B, A). The A controls transparency and is optional, so 64 × 3 would work.

Quiver with variable direction, position, and colour

Here we are animating 30 arrows. Qx and Qy are the locations of the base of the arrows for all frames. Their size is (N_f, 30).  U and V, the arrows’ directions, are also this size. C, the colours of the arrows, is (N_f, 4), but for each frame we take a subset of size (30, 4).

def animate(i):
    s = np.s_[i, :]
    qax.set_offsets(np.c_[Qx[s].flatten(), Qy[s].flatten()])
    qax.set_UVC(U[s], V[s])
    Cidx = np.mod(np.r_[i:i+U.shape[1]], Nframes)
    qax.set_facecolor(C[Cidx, :])

• set_offsets works as it does for the scatter. Since we have 30 arrows, it accepts an array of size (30, 2).
• set_UVC accepts the arrow vectors as two 1D arrays, rather than combined.
• set_facecolor accepts an array of size (30, 4), one RGBA colour for each arrow. As above, the A is optional.

The first line within the function simplifies the slicing of arrays in the following lines. The fourth line involving the modulo division (np.mod) is only there to help cycle through the colourmap.

Line plots in a polar projection

Because Matplotlib does all the hard work behind the scenes with respect to graph projections, animating a polar plot is no different to animating a line plot. Here we update two curves. theta1 and theta2 are 2D arrays of size (N_f, 100), where 100 is the number of elements comprising a curve for a given frame.

def animate(i):
    curve1.set_ydata(theta1[i, :])
    curve2.set_ydata(theta2[i, :])

Filled area plots with variable colour

Filled area plots are a little more challenging as they are patches, rather than points or lines. Matplotlib’s treatment of patches makes sense, but it isn’t intuitive at the outset which properties control what.

fill is an object created by the fill_between function. Y1 and Y2, the upper and lower bounds of the fill, are each arrays of size (N_f, N_x), where N_x is the number of elements making up each curve in a single frame.

def animate(i):
    path = fill.get_paths()[0]
    verts = path.vertices
    verts[1:Nx+1, 1] = Y1[i, :]
    verts[Nx+2:-1, 1] = Y2[i, :][::-1]
    fill.set_color(C[i, :])

• get_paths()[0] extracts the one and only path. The output of get_paths() is a list, even if only one element long.
• We rename path.vertices for clarity.
• The elements 1:Nx+1 contain the y coordinates of the upper curve.
• The elements Nx+2:-1 contain the y coordinates of the lower curve, but right-to-left, hence the need to reverse the array using [::-1].
• set_color works as it does for examples above.

Careful inspection of the third and fourth lines of the function indicates that the 0, Nx+1, and -1 elements of verts are not being updated. It’s not clear what these elements do.

Three-dimensional contour plots

Contour plots are treated different from the other examples. In a way, they are simpler, because for each frame, we simply delete the current contour object and then plot a new one. No need to figure out what set_... method to use. Although any animation could be created this way, the trade-off is often a noticeable slow down in the speed that the animation renders and unnecessary repetition in the Python code.

Here X and Y  are 2D arrays of size (20, 30) and Z, the height of the contour, is a 3D array of size (N_f, 20, 30).

def animate(i):
    ax.collections = []
    ax.plot_surface(X, Y, Z[i, :, :], **contour_opts)				data-mce-type="bookmark" 				id="mce_SELREST_start" 				data-mce-style="overflow:hidden;line-height:0" 				style="overflow:hidden;line-height:0" 			></span>

• Delete the existing contour by emptying the list of collections, which is where Matplotlib stores information about the contour object.
• Replot the contour using plot_surface with **contour_opts being a dictionary of keyword arguments.

With judicious use of ax.collections, it is possible to retain some components of a plot while removing others. Here that’s unnecessary since we’re plotting only a single thing.