Auxillary tutorial 5: Saving a movie using matplotlib

This auxillary tutorial was generated from an IPython notebook. You can download the notebook here.

In [1]: 
# Stuff we always import
from __future__ import print_function, division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

# Inline for the purposes of making this document
%matplotlib inline

# And an ODE integrator for our fun curve
import scipy.integrate

We often want to store movies either of plots or of a series of images we analyze. The XYTStack class of jb_utils.py has a method show_movie for displaying movies, which is backend dependent. Here, we will show three ways of making movies. First, we can store each successive frame of a movie as PNGs and then use some other software (such as Quicktime, ffmpeg, or Fiji to assemble them into a movie. Second, we can use the matplotlib.animate moduile, which uses ffmpeg under the hood to convert the figures into images. Finally, we will use HTML5's capabilities for embedded videa. This last method is browser dependent; not all browsers support it. I tested it with Safari 8 and Chrom 38.

As in our tutorial for saving a figure, we will use the Lorentz attractor to make our movie.

The Lorentz attractor is described by a system of ordinary differential equations

\begin{align} \frac{\mathrm{d}x}{\mathrm{d}t} & = s(y-x) \\ \frac{\mathrm{d}y}{\mathrm{d}t} & = x(p-z) - y \\ \frac{\mathrm{d}z}{\mathrm{d}t} & = xy - bz. \end{align}

For certain sets of parameters $(s, p, b)$, the solution to the Lorentz attractor exhibits chaotic trajectories, which is why it is fun. We will plot a projection of the trajectory onto the $x$-$y$ plane.

In [6]: 
"""
This is a bunch of code to generate a Lorentz attractor that you can skip
and go on to saving the figure.
"""
def lorentz_attractor(r, t, p):
    """
    Right hand side of system of ODEs describing Lorentz attractor.
    x' = s * (y - x)
    y' = x * (p - z) - y
    z' = x * y - b * z
    
    The input r is (x, y, z) and p is (s, p, b).  
    
    t is time and must be present for odeint to work, but not explicit
    in computation of the Lorentz attractor ODE.
    """
    x, y, z = r
    s, p, b = p
    return np.array([s * (y - x), x * (p - z) - y, x * y - b * z])

# Parameters to use
p = np.array([10.0, 28.0, 8.0 / 3.0])

# Initial condition
r0 = np.array([0.1, 0.0, 0.0])

# Time points to sample
t = np.linspace(0.0, 40.0, 500)

# Use scipy.integrate.odeint to integrate Lorentz attractor
r = scipy.integrate.odeint(lorentz_attractor, r0, t, args=(p,))

# Unpack results into x, y, z.
x, y, z = r.transpose()

Now that we have our trajectory, we can make our movies. We'll start by individually saving PNGs for each and then assembling them.

In [21]: 
# Set up figure
fig = plt.figure()
ax = plt.axes(xlim=(-20, 20), ylim=(-2, 50))
ax.axis('off')

# Populate figure with empty plots (will be dot and trail)
plots = [ax.plot([], [], 'r.', markersize=6, zorder=1)[0],
         ax.plot([], [], 'k-', lw=0.5, zorder=0)[0]]

# Loop through trajectory and plot x,z position and save figure
for i in range(len(t)):
    plots[0].set_data(x[i], z[i])
    plots[1].set_data(x[:i], z[:i])
    fig.savefig('frame_%04d.png' % i, dpi=100)

After doing this, we will have a set of PNGs named frame_0001.png, frame_0002.png, etc., that we can put together to make a movie using some other software.

As an alternative, we can use matplotlib's animate function to do this automatically. You will need ffmpeg installed for this to work.

In [24]: 
# Close the previous figure to start over
plt.close('all')

# Set up figure
fig = plt.figure()
ax = plt.axes(xlim=(-20, 20), ylim=(-2, 50))
ax.axis('off')

# Populate figure with empty plots (will be dot and trail)
plots = [ax.plot([], [], 'r.', markersize=6, zorder=1)[0],
         ax.plot([], [], 'k-', lw=0.5, zorder=0)[0]]

# Initialize plotting
def init():
    plots[0].set_data([],[])
    plots[1].set_data([],[])
    return plots

# This paints in the animation
def animate(i):
    plots[0].set_data(x[i], z[i])
    plots[1].set_data(x[:i], z[:i])
    return plots

# call the animator.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=len(t), interval=20, blit=True)

# Save the animation
anim.save('lorentz.mp4', fps=24, dpi=100, 
          extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p'])

# Close the figure
plt.close()

Finally, we can embed the encoded MP4 movie in HTML.

In [25]: 
# Import modules for encoding and embedding movie
from IPython.display import HTML
from base64 import b64encode

# Encode/embed the animation
video = open("lorentz.mp4", "rb").read()
video_encoded = b64encode(video)
video_tag = """
<video controls>
<source alt="G-M video" src="data:video/x-m4v;base64,{0}" type="video/mp4">
Your browser does not support MP4 H.264 video.
</video>
""".format(video_encoded)
HTML(data=video_tag)
Out[25]: