matplotlib · ghost · Jun 8, 2015 · Jun 8, 2015 · Jun 8, 2015 · Sep 25, 2015
diff --git a/doc/mpl_toolkits/index.rst b/doc/mpl_toolkits/index.rst
@@ -198,3 +198,10 @@ external matplotlib backend uses iTerm2 nightly build inline image display
 feature.
 
 .. image:: /_static/matplotlib_iterm2_demo.png
+
+mpl_parsave
+===========
+
+Provides a class, Parsave, to fascilitate recording of animations in parallel.
+That is, an animation is broken into blocks, blocks are recorded in parallel
+using multiprocessing, then stitch the blocks into a final movie.
diff --git a/lib/mpl_toolkits/mpl_parsave/__init__.py b/lib/mpl_toolkits/mpl_parsave/__init__.py
@@ -0,0 +1,2 @@
+
+__all__ = ["mplparsave"]
diff --git a/lib/mpl_toolkits/mpl_parsave/examples.py b/lib/mpl_toolkits/mpl_parsave/examples.py
@@ -0,0 +1,112 @@
+from numpy import sin, cos, pi, array
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.integrate as integrate
+import matplotlib.animation as animation
+
+import mplparsave
+
+# THIS CODE COMES FROM
+# matplotlib/examples/animation/dynamic_pendulum_animated.py:
+
+G =  9.8 # acceleration due to gravity, in m/s^2
+L1 = 1.0 # length of pendulum 1 in m
+L2 = 1.0 # length of pendulum 2 in m
+M1 = 1.0 # mass of pendulum 1 in kg
+M2 = 1.0 # mass of pendulum 2 in kg
+
+
+def derivs(state, t):
+
+    dydx = np.zeros_like(state)
+    dydx[0] = state[1]
+
+    del_ = state[2]-state[0]
+    den1 = (M1+M2)*L1 - M2*L1*cos(del_)*cos(del_)
+    dydx[1] = (M2*L1*state[1]*state[1]*sin(del_)*cos(del_)
+               + M2*G*sin(state[2])*cos(del_) + M2*L2*state[3]*state[3]*sin(del_)
+               - (M1+M2)*G*sin(state[0]))/den1
+
+    dydx[2] = state[3]
+
+    den2 = (L2/L1)*den1
+    dydx[3] = (-M2*L2*state[3]*state[3]*sin(del_)*cos(del_)
+               + (M1+M2)*G*sin(state[0])*cos(del_)
+               - (M1+M2)*L1*state[1]*state[1]*sin(del_)
+               - (M1+M2)*G*sin(state[2]))/den2
+
+    return dydx
+
+# create a time array from 0..100 sampled at 0.05 second steps
+dt = 0.05
+t = np.arange(0.0, 40, dt)
+
+# th1 and th2 are the initial angles (degrees)
+# w10 and w20 are the initial angular velocities (degrees per second)
+th1 = 120.0
+w1 = 0.0
+th2 = -10.0
+w2 = 0.0
+
+rad = pi/180
+
+# initial state
+state = np.array([th1, w1, th2, w2])*pi/180.
+
+# integrate your ODE using scipy.integrate.
+y = integrate.odeint(derivs, state, t)
+
+x1 = L1*sin(y[:,0])
+y1 = -L1*cos(y[:,0])
+
+x2 = L2*sin(y[:,2]) + x1
+y2 = -L2*cos(y[:,2]) + y1
+
+fig = plt.figure()
+ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 2))
+ax.grid()
+
+line, = ax.plot([], [], 'o-', lw=2)
+time_template = 'time = %.1fs'
+time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
+
+def init():
+    line.set_data([], [])
+    time_text.set_text('')
+    return line, time_text
+
+def animate(i):
+    thisx = [0, x1[i], x2[i]]
+    thisy = [0, y1[i], y2[i]]
+
+    line.set_data(thisx, thisy)
+    time_text.set_text(time_template%(i*dt))
+    return line, time_text
+
+############## THIS CODE IS TO DEMONSTRATE HOW MPLPARSAVE WORKS ############
+
+def parsave():
+    # The first half of the frames...
+    block1=np.arange(1, len(y)/2)
+    # The second half of the frames...
+    block2=np.arange(len(y)/2, len(y))
+    # Two blocks in total will tell the recorder to run two processes in
+    # parallel, one for each of the blocks, and stitch the two at the end.
+    # More blocks will trigger more processes. One shouldn't run more than
+    # the number of available cores.
+    blocks=[block1, block2]
+
+    # Set up the writer class (matplotlib proper). We use 'ffmpeg', but
+    # any of the supported writers can be used.
+    Writer = animation.writers['ffmpeg']
+    writer = Writer(fps=30, bitrate=1800)
+
+    # Record and stitch using ffmpeg as stitcher...
+    mplparsave.Parsave.record("ffmpeg-pendulum.mp4", fig, animate, init,
+                              blocks, writer, mplparsave.Stitcher('ffmpeg'),
+                              interval=35)
+
+    # Record and stitch using mencoder as stitcher...
+    mplparsave.Parsave.record("mencoder-pendulum.mp4", fig, animate, init,
+                              blocks, writer, mplparsave.Stitcher('mencoder'),
+                              interval=35)
diff --git a/lib/mpl_toolkits/mpl_parsave/mplparsave.py b/lib/mpl_toolkits/mpl_parsave/mplparsave.py
@@ -0,0 +1,184 @@
+"""
+Speed up writing matplotlib animations to a movie file by using multiple
+processes.
+
+Contains the Parsave class. Parsave.record method fascilitates recording
+of animations in parallel using multiple processes.
+"""
+
+import os
+from matplotlib.animation import FuncAnimation as fanim
+from multiprocessing import Process
+import subprocess
+import time
+
+class Parsave(object):
+  """
+  Parsave: class containing static methods to fascilitate parallel recording.
+  """
+  def __init__(self):
+    pass
+
+  @staticmethod
+  def record(fname, fig, anim_func, init_func, blocks,
+             writerClass, stitcherClass, keep=False, **kwargs):
+
+"""
+    record(fname, fig, anim_func, init_func, blocks, writerClass,
+           stitcherClass, keep=False, **kwargs)
+
+    Record the movie. The parameters are
+
+    fname: output file name.
+
+    init: the function that draws background of each frame.
+
+    fig: figure where animation will be drawn.
+
+    anim_func: the function to be animated.
+
+    blocks: an array of the form [block_1, block_2, ..., block_n]
+    where for j = 1, ..., n, block_j is the (array of) frames to be recorded
+    by the jth thread using writerClass, which at the end is stitched together
+    into a single movie using stitcherClass.
+
+    writerClass: the writer responsible for writing the frames
+    (an instance of the matplotlib.animation.MovieWriter).
+
+    stitcherClass: the stitcher responsible for stitching the movies recorded
+    by the writerClass into a single movie.
+
+    The number of threads that will be run is equal to n
+    (the number of blocks as above).
+
+    keep: True or False; whether to keep the parts that are then stitched
+    into the final movie, or delete them after stitching the final movie.
+
+    **kwargs are the keyword arguments to be passed to anim_func.
+"""
+
+    num_jobs=len(blocks)
+    jobs=[0]*num_jobs
+    names=[0]*num_jobs
+
+    # Record movies.
+    for j in range(num_jobs):
+      names[j]=(('%d'+'-'+'%f'+'.'+fname.split('.')[-1])) % \
+                (j, float(time.time()))
+      jobs[j]=Process(target=Parsave.__parallel_save,
+                      args=(names[j], fig, init_func, anim_func, blocks[j],
+                            writerClass), kwargs=kwargs)
+      jobs[j].start()
+
+    for job in jobs:
+      job.join()
+
+    # Now stitch the recorded movies together.
+    if num_jobs==1:
+      os.rename(names[0], fname)
+    else:
+      stitcherClass.stitch(fname, names)
+      if keep is False:
+        for name in names:
+          os.remove(name)
+
+  # The helpber for record(). 'frames' is a block from 'blocks' passed to
+  # record().
+  @staticmethod
+  def __parallel_save(fname, fig, init, anim_func,
+                      frames, writerClass, **kwargs):
+    anm=fanim(fig, anim_func, init_func=init, frames=frames, **kwargs)
+    anm.save(fname, writer=writerClass)
+# ========================================================================#
+
+
+
+
+
+
+
+# === STITCHER WRAPPER ===================================================#
+# Wrapper for the stitcher classes
+# ========================================================================#
+supported_stitchers=['ffmpeg',
+                     'mencoder']
+
+class Stitcher(object):
+  # name: the name of the class to use, e.g. 'mencoder'.
+  #
+  # args: arguments for the stitcher.
+  def __init__(self, name, args=None):
+    if name=='ffmpeg':
+      self.stitcher=FFMpeg(args)
+    elif name=='mencoder':
+      self.stitcher=Mencoder(args)
+    else:
+      error='Unsupported stitcher. Supported stitchers: '
+      for s in supported_stitchers:
+        error+=(s+', ')
+      error=error[:-2]+'.'
+      raise ValueError(error)
+
+  # fname: the output file name.
+  # fnames = [fname_1, fname_2, ..., fname_n] the file names of the movies
+  # to stitch.
+  def stitch(self, fname, fnames):
+    self.stitcher.stitch(fname, fnames)
+# ========================================================================#
+
+
+
+
+
+
+
+# === STITCHERS ==========================================================#
+# Currently supported:
+#  - mencoder
+#  - ffmpeg
+#
+#
+# === MENCODER STITCHER ==================================================#
+# Stitcher based on 'mencoder'.
+# ========================================================================#
+class Mencoder(object):
+  def __init__(self, args=None):
+    self.args=args
+
+  def stitch(self, fname, fnames):
+    if self.args is None:
+      self.args=['-ovc', 'copy', '-idx', '-o']+[fname]+fnames
+
+    subprocess.check_call(['mencoder']+self.args,
+                          stdout=open(os.devnull, 'w'),
+                          stderr=subprocess.STDOUT)
+# ========================================================================#
+
+
+
+
+
+
+
+# === FFMPEG STITCHER ====================================================#
+# Stitcher based on 'ffmpeg'.
+# ========================================================================#
+class FFMpeg(object):
+  def __init__(self, args=None):
+    self.args=args
+
+  def stitch(self, fname, fnames):
+    input_file=(('%s%f'+'.txt')%('input', float(time.time())))
+    f=open(input_file, 'w')
+    for name in fnames:
+      f.write('file '+"'"+name+"'"+'\n')
+    f.close()
+    args=['-f', 'concat', '-i', input_file, '-codec', 'copy']+[fname]
+    if self.args is not None:
+      args=self.args+args
+
+    subprocess.check_call(['ffmpeg']+args,
+                          stdout=open(os.devnull, 'w'),
+                          stderr=subprocess.STDOUT)
+    os.remove(input_file)
+# ========================================================================#
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@

		__all__ = ["mplparsave"]
Copy link Member tacaswell Nov 17, 2015 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. I would get rid of this package and just make a `parsave` module under `mpl_toolkits`