'''Compute figure size so that pixels are a multiple of n

# this maybe simplified if / when we adopt consistent rounding for

# pixel size across the whole library

def correct_roundoff(x, dpi, n):

if int(x*dpi) % n != 0:

if int(np.nextafter(x, np.inf)*dpi) % n == 0:

x = np.nextafter(x, np.inf)

elif int(np.nextafter(x, -np.inf)*dpi) % n == 0:

x = np.nextafter(x, -np.inf)

return x

wnew = int(w * dpi / n) * n / dpi

hnew = int(h * dpi / n) * n / dpi

'''Registry of available writer classes by human readable name.'''

def __init__(self):

self.avail = dict()

self._registered = dict()

self._dirty = False

def set_dirty(self):

"""Sets a flag to re-setup the writers"""

self._dirty = True

# Returns a decorator that can be used on classes to register them under

# a name. As in:

# @register('foo')

# class Foo:

# pass

def register(self, name):

def wrapper(writerClass):

self._registered[name] = writerClass

if writerClass.isAvailable():

self.avail[name] = writerClass

return writerClass

return wrapper

def ensure_not_dirty(self):

"""If dirty, reasks the writers if they are available"""

if self._dirty:

self.reset_available_writers()

def reset_available_writers(self):

"""Reset the available state of all registered writers"""

self.avail = {}

for name, writerClass in self._registered.items():

if writerClass.isAvailable():

self.avail[name] = writerClass

self._dirty = False

def list(self):

''' Get a list of available MovieWriters.'''

self.ensure_not_dirty()

return list(self.avail)

def is_available(self, name):

'''Check if given writer is available by name.

self.ensure_not_dirty()

return name in self.avail

def __getitem__(self, name):

self.ensure_not_dirty()

if not self.avail:

raise RuntimeError("No MovieWriters available!")

'''

@abc.abstractmethod

def setup(self, fig, outfile, dpi=None):

'''

@abc.abstractmethod

def grab_frame(self, **savefig_kwargs):

'''

@abc.abstractmethod

def finish(self):

'Finish any processing for writing the movie.'

@contextlib.contextmanager

def saving(self, fig, outfile, dpi, *args, **kwargs):

'''

# This particular sequence is what contextlib.contextmanager wants

self.setup(fig, outfile, dpi, *args, **kwargs)

try:

yield self

finally:

'''Base class for writing movies.

def __init__(self, fps=5, codec=None, bitrate=None, extra_args=None,

metadata=None):

'''MovieWriter

self.fps = fps

self.frame_format = 'rgba'

if codec is None:

self.codec = rcParams['animation.codec']

else:

self.codec = codec

if bitrate is None:

self.bitrate = rcParams['animation.bitrate']

else:

self.bitrate = bitrate

if extra_args is None:

self.extra_args = list(rcParams[self.args_key])

else:

self.extra_args = extra_args

if metadata is None:

self.metadata = dict()

else:

self.metadata = metadata

@property

def frame_size(self):

'A tuple (width,height) in pixels of a movie frame.'

w, h = self.fig.get_size_inches()

return int(w * self.dpi), int(h * self.dpi)

def _adjust_frame_size(self):

if self.codec == 'h264':

wo, ho = self.fig.get_size_inches()

w, h = adjusted_figsize(wo, ho, self.dpi, 2)

if not (wo, ho) == (w, h):

self.fig.set_size_inches(w, h, forward=True)

verbose.report('figure size (inches) has been adjusted '

'from %s x %s to %s x %s' % (wo, ho, w, h),

level='helpful')

else:

w, h = self.fig.get_size_inches()

verbose.report('frame size in pixels is %s x %s' % self.frame_size,

level='debug')

return w, h

def setup(self, fig, outfile, dpi=None):

'''

self.outfile = outfile

self.fig = fig

if dpi is None:

dpi = self.fig.dpi

self.dpi = dpi

self._w, self._h = self._adjust_frame_size()

# Run here so that grab_frame() can write the data to a pipe. This

# eliminates the need for temp files.

self._run()

def _run(self):

# Uses subprocess to call the program for assembling frames into a

# movie file. *args* returns the sequence of command line arguments

# from a few configuration options.

command = self._args()

if verbose.ge('debug'):

output = sys.stdout

else:

output = subprocess.PIPE

verbose.report('MovieWriter.run: running command: %s' %

' '.join(command))

self._proc = subprocess.Popen(command, shell=False,

stdout=output, stderr=output,

stdin=subprocess.PIPE,

creationflags=subprocess_creation_flags)

def finish(self):

'Finish any processing for writing the movie.'

self.cleanup()

def grab_frame(self, **savefig_kwargs):

'''

verbose.report('MovieWriter.grab_frame: Grabbing frame.',

level='debug')

try:

# re-adjust the figure size in case it has been changed by the

# user. We must ensure that every frame is the same size or

# the movie will not save correctly.

self.fig.set_size_inches(self._w, self._h)

# Tell the figure to save its data to the sink, using the

# frame format and dpi.

self.fig.savefig(self._frame_sink(), format=self.frame_format,

dpi=self.dpi, **savefig_kwargs)

except (RuntimeError, IOError) as e:

out, err = self._proc.communicate()

verbose.report('MovieWriter -- Error '

'running proc:\n%s\n%s' % (out,

err), level='helpful')

raise IOError('Error saving animation to file (cause: {0}) '

'Stdout: {1} StdError: {2}. It may help to re-run '

'with --verbose-debug.'.format(e, out, err))

def _frame_sink(self):

'Returns the place to which frames should be written.'

return self._proc.stdin

def _args(self):

'Assemble list of utility-specific command-line arguments.'

return NotImplementedError("args needs to be implemented by subclass.")

def cleanup(self):

'Clean-up and collect the process used to write the movie file.'

out, err = self._proc.communicate()

self._frame_sink().close()

verbose.report('MovieWriter -- '

'Command stdout:\n%s' % out, level='debug')

verbose.report('MovieWriter -- '

'Command stderr:\n%s' % err, level='debug')

@classmethod

def bin_path(cls):

'''

return str(rcParams[cls.exec_key])

@classmethod

def isAvailable(cls):

'''

bin_path = cls.bin_path()

if not bin_path:

return False

try:

p = subprocess.Popen(

bin_path,

shell=False,

stdout=subprocess.PIPE,

stderr=subprocess.PIPE,

creationflags=subprocess_creation_flags)

return cls._handle_subprocess(p)

except OSError:

return False

@classmethod

def _handle_subprocess(cls, process):

process.communicate()

'''`MovieWriter` for writing to individual files and stitching at the end.

def __init__(self, *args, **kwargs):

MovieWriter.__init__(self, *args, **kwargs)

self.frame_format = rcParams['animation.frame_format']

def setup(self, fig, outfile, dpi=None, frame_prefix='_tmp',

clear_temp=True):

'''Perform setup for writing the movie file.

self.fig = fig

self.outfile = outfile

if dpi is None:

dpi = self.fig.dpi

self.dpi = dpi

self._adjust_frame_size()

self.clear_temp = clear_temp

self.temp_prefix = frame_prefix

self._frame_counter = 0 # used for generating sequential file names

self._temp_names = list()

self.fname_format_str = '%s%%07d.%s'

@property

def frame_format(self):

'''

return self._frame_format

@frame_format.setter

def frame_format(self, frame_format):

if frame_format in self.supported_formats:

self._frame_format = frame_format

else:

self._frame_format = self.supported_formats[0]

def _base_temp_name(self):

# Generates a template name (without number) given the frame format

# for extension and the prefix.

return self.fname_format_str % (self.temp_prefix, self.frame_format)

def _frame_sink(self):

# Creates a filename for saving using the basename and the current

# counter.

fname = self._base_temp_name() % self._frame_counter

# Save the filename so we can delete it later if necessary

self._temp_names.append(fname)

verbose.report(

'FileMovieWriter.frame_sink: saving frame %d to fname=%s' %

(self._frame_counter, fname),

level='debug')

self._frame_counter += 1 # Ensures each created name is 'unique'

# This file returned here will be closed once it's used by savefig()

# because it will no longer be referenced and will be gc-ed.

return open(fname, 'wb')

def grab_frame(self, **savefig_kwargs):

'''

# Overloaded to explicitly close temp file.

verbose.report('MovieWriter.grab_frame: Grabbing frame.',

level='debug')

try:

# Tell the figure to save its data to the sink, using the

# frame format and dpi.

with self._frame_sink() as myframesink:

self.fig.savefig(myframesink, format=self.frame_format,

dpi=self.dpi, **savefig_kwargs)

except RuntimeError:

out, err = self._proc.communicate()

verbose.report('MovieWriter -- Error '

'running proc:\n%s\n%s' % (out,

err), level='helpful')

raise

def finish(self):

# Call run here now that all frame grabbing is done. All temp files

# are available to be assembled.

self._run()

MovieWriter.finish(self) # Will call clean-up

# Check error code for creating file here, since we just run

# the process here, rather than having an open pipe.

if self._proc.returncode:

try:

stdout = [s.decode() for s in self._proc._stdout_buff]

stderr = [s.decode() for s in self._proc._stderr_buff]

verbose.report("MovieWriter.finish: stdout: %s" % stdout,

level='helpful')

verbose.report("MovieWriter.finish: stderr: %s" % stderr,

level='helpful')

except Exception as e:

pass

msg = ('Error creating movie, return code: ' +

str(self._proc.returncode) +

' Try setting mpl.verbose.set_level("helpful")')

raise RuntimeError(msg)

def cleanup(self):

MovieWriter.cleanup(self)

# Delete temporary files

if self.clear_temp:

verbose.report(

'MovieWriter: clearing temporary fnames=%s' %

str(self._temp_names),

level='debug')

for fname in self._temp_names:

'''Mixin class for FFMpeg output.

exec_key = 'animation.ffmpeg_path'

args_key = 'animation.ffmpeg_args'

@property

def output_args(self):

args = ['-vcodec', self.codec]

# For h264, the default format is yuv444p, which is not compatible

# with quicktime (and others). Specifying yuv420p fixes playback on

# iOS,as well as HTML5 video in firefox and safari (on both Win and

# OSX). Also fixes internet explorer. This is as of 2015/10/29.

if self.codec == 'h264' and '-pix_fmt' not in self.extra_args:

args.extend(['-pix_fmt', 'yuv420p'])

# The %dk adds 'k' as a suffix so that ffmpeg treats our bitrate as in

# kbps

if self.bitrate > 0:

args.extend(['-b', '%dk' % self.bitrate])

if self.extra_args:

args.extend(self.extra_args)

for k, v in six.iteritems(self.metadata):

args.extend(['-metadata', '%s=%s' % (k, v)])

return args + ['-y', self.outfile]

@classmethod

def _handle_subprocess(cls, process):

_, err = process.communicate()

# Ubuntu 12.04 ships a broken ffmpeg binary which we shouldn't use

# NOTE : when removed, remove the same method in AVConvBase.

if 'Libav' in err.decode():

return False

'''Pipe-based ffmpeg writer.

def _args(self):

# Returns the command line parameters for subprocess to use

# ffmpeg to create a movie using a pipe.

args = [self.bin_path(), '-f', 'rawvideo', '-vcodec', 'rawvideo',

'-s', '%dx%d' % self.frame_size, '-pix_fmt', self.frame_format,

'-r', str(self.fps)]

# Logging is quieted because subprocess.PIPE has limited buffer size.

if not verbose.ge('debug'):

args += ['-loglevel', 'quiet']

args += ['-i', 'pipe:'] + self.output_args

'''File-based ffmpeg writer.

supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',

'pbm', 'raw', 'rgba']

def _args(self):

# Returns the command line parameters for subprocess to use

# ffmpeg to create a movie using a collection of temp images

return [self.bin_path(), '-r', str(self.fps),

'-i', self._base_temp_name(),

'''Mixin class for avconv output.

exec_key = 'animation.avconv_path'

args_key = 'animation.avconv_args'

# NOTE : should be removed when the same method is removed in FFMpegBase.

@classmethod

def _handle_subprocess(cls, process):

'''Pipe-based avconv writer.

'''File-based avconv writer.

exec_key = 'animation.mencoder_path'

args_key = 'animation.mencoder_args'

# Mencoder only allows certain keys, other ones cause the program

# to fail.

allowed_metadata = ['name', 'artist', 'genre', 'subject', 'copyright',

'srcform', 'comment']

# Mencoder mandates using name, but 'title' works better with ffmpeg.

# If we find it, just put it's value into name

def _remap_metadata(self):

if 'title' in self.metadata:

self.metadata['name'] = self.metadata['title']

@property

def output_args(self):

self._remap_metadata()

lavcopts = {'vcodec': self.codec}

if self.bitrate > 0:

lavcopts.update(vbitrate=self.bitrate)

args = ['-o', self.outfile, '-ovc', 'lavc', '-lavcopts',

':'.join(itertools.starmap('{0}={1}'.format,

lavcopts.items()))]

if self.extra_args:

args.extend(self.extra_args)

if self.metadata:

args.extend(['-info', ':'.join('%s=%s' % (k, v)

for k, v in six.iteritems(self.metadata)

if k in self.allowed_metadata)])

@deprecated('2.0', message=mencoder_dep)

def __init__(self, *args, **kwargs):

with rc_context(rc={'animation.codec': 'mpeg4'}):

super(MencoderWriter, self).__init__(*args, **kwargs)

def _args(self):

# Returns the command line parameters for subprocess to use

# mencoder to create a movie

return [self.bin_path(), '-', '-demuxer', 'rawvideo', '-rawvideo',

('w=%i:h=%i:' % self.frame_size +

'fps=%i:format=%s' % (self.fps,

supported_formats = ['png', 'jpeg', 'tga', 'sgi']

@deprecated('2.0', message=mencoder_dep)

def __init__(self, *args, **kwargs):

with rc_context(rc={'animation.codec': 'mpeg4'}):

super(MencoderFileWriter, self).__init__(*args, **kwargs)

def _args(self):

# Returns the command line parameters for subprocess to use

# mencoder to create a movie

return [self.bin_path(),

'mf://%s*.%s' % (self.temp_prefix, self.frame_format),

'-frames', str(self._frame_counter), '-mf',

'type=%s:fps=%d' % (self.frame_format,

'''Mixin class for ImageMagick output.

exec_key = 'animation.convert_path'

args_key = 'animation.convert_args'

@property

def delay(self):

return 100. / self.fps

@property

def output_args(self):

return [self.outfile]

@classmethod

def _init_from_registry(cls):

if sys.platform != 'win32' or rcParams[cls.exec_key] != 'convert':

return

from six.moves import winreg

for flag in (0, winreg.KEY_WOW64_32KEY, winreg.KEY_WOW64_64KEY):

try:

hkey = winreg.OpenKeyEx(winreg.HKEY_LOCAL_MACHINE,

'Software\\Imagemagick\\Current',

0, winreg.KEY_QUERY_VALUE | flag)

binpath = winreg.QueryValueEx(hkey, 'BinPath')[0]

winreg.CloseKey(hkey)

binpath += '\\convert.exe'

break

except Exception:

binpath = ''

rcParams[cls.exec_key] = rcParamsDefault[cls.exec_key] = binpath

@classmethod

def isAvailable(cls):

'''

bin_path = cls.bin_path()

if bin_path == "convert":

cls._init_from_registry()

'''Pipe-based animated gif.

def _args(self):

return ([self.bin_path(),

'-size', '%ix%i' % self.frame_size, '-depth', '8',

'-delay', str(self.delay), '-loop', '0',

'%s:-' % self.frame_format]

'''File-based animated gif writer.

supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',

'pbm', 'raw', 'rgba']

def _args(self):

return ([self.bin_path(), '-delay', str(self.delay), '-loop', '0',

'%s*.%s' % (self.temp_prefix, self.frame_format)]

'''This class wraps the creation of an animation using matplotlib.

def __init__(self, fig, event_source=None, blit=False):

self._fig = fig

# Disables blitting for backends that don't support it. This

# allows users to request it if available, but still have a

# fallback that works if it is not.

self._blit = blit and fig.canvas.supports_blit

# These are the basics of the animation. The frame sequence represents

# information for each frame of the animation and depends on how the

# drawing is handled by the subclasses. The event source fires events

# that cause the frame sequence to be iterated.

self.frame_seq = self.new_frame_seq()

self.event_source = event_source

# Instead of starting the event source now, we connect to the figure's

# draw_event, so that we only start once the figure has been drawn.

self._first_draw_id = fig.canvas.mpl_connect('draw_event', self._start)

# Connect to the figure's close_event so that we don't continue to

# fire events and try to draw to a deleted figure.

self._close_id = self._fig.canvas.mpl_connect('close_event',

self._stop)

if self._blit:

self._setup_blit()

def _start(self, *args):

'''

# First disconnect our draw event handler

self._fig.canvas.mpl_disconnect(self._first_draw_id)

self._first_draw_id = None # So we can check on save

# Now do any initial draw

self._init_draw()

# Add our callback for stepping the animation and

# actually start the event_source.

self.event_source.add_callback(self._step)

self.event_source.start()

def _stop(self, *args):

# On stop we disconnect all of our events.

if self._blit:

self._fig.canvas.mpl_disconnect(self._resize_id)

self._fig.canvas.mpl_disconnect(self._close_id)

self.event_source.remove_callback(self._step)

self.event_source = None

def save(self, filename, writer=None, fps=None, dpi=None, codec=None,

bitrate=None, extra_args=None, metadata=None, extra_anim=None,

savefig_kwargs=None):

'''Saves a movie file by drawing every frame.