# TODO: # * Loop Delay is broken on GTKAgg. This is because source_remove() is not # working as we want. PyGTK bug? # * Documentation -- this will need a new section of the User's Guide. # Both for Animations and just timers. # - Also need to update http://www.scipy.org/Cookbook/Matplotlib/Animations # * Blit # * Currently broken with Qt4 for widgets that don't start on screen # * Still a few edge cases that aren't working correctly # * Can this integrate better with existing matplotlib animation artist flag? # - If animated removes from default draw(), perhaps we could use this to # simplify initial draw. # * Example # * Frameless animation - pure procedural with no loop # * Need example that uses something like inotify or subprocess # * Complex syncing examples # * Movies # * Library to make movies? # * RC parameter for config? # * Can blit be enabled for movies? # * Need to consider event sources to allow clicking through multiple figures import itertools from matplotlib.cbook import iterable from matplotlib import verbose class Animation(object): ''' This class wraps the creation of an animation using matplotlib. It is only a base class which should be subclassed to provide needed behavior. *fig* is the figure object that is used to get draw, resize, and any other needed events. *event_source* is a class that can run a callback when desired events are generated, as well as be stopped and started. Examples include timers (see :class:`TimedAnimation`) and file system notifications. *blit* is a boolean that controls whether blitting is used to optimize drawing. ''' def __init__(self, fig, event_source=None, blit=False): self._fig = fig self._blit = 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 # Clear the initial frame self._init_draw() # 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 blit: self._setup_blit() def _start(self, *args): ''' Starts interactive animation. Adds the draw frame command to the GUI handler, calls show to start the event loop. ''' # On start, we add our callback for stepping the animation and # actually start the event_source. We also disconnect _start # from the draw_events self.event_source.add_callback(self._step) self.event_source.start() self._fig.canvas.mpl_disconnect(self._first_draw_id) self._first_draw_id = None # So we can check on save 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, fps=5, codec='mpeg4', clear_temp=True, frame_prefix='_tmp'): ''' Saves a movie file by drawing every frame. *filename* is the output filename, eg :file:`mymovie.mp4` *fps* is the frames per second in the movie *codec* is the codec to be used,if it is supported by the output method. *clear_temp* specifies whether the temporary image files should be deleted. *frame_prefix* gives the prefix that should be used for individual image files. This prefix will have a frame number (i.e. 0001) appended when saving individual frames. ''' # Need to disconnect the first draw callback, since we'll be doing # draws. Otherwise, we'll end up starting the animation. if self._first_draw_id is not None: self._fig.canvas.mpl_disconnect(self._first_draw_id) reconnect_first_draw = True else: reconnect_first_draw = False fnames = [] # Create a new sequence of frames for saved data. This is different # from new_frame_seq() to give the ability to save 'live' generated # frame information to be saved later. # TODO: Right now, after closing the figure, saving a movie won't # work since GUI widgets are gone. Either need to remove extra code # to allow for this non-existant use case or find a way to make it work. for idx,data in enumerate(self.new_saved_frame_seq()): #TODO: Need to see if turning off blit is really necessary self._draw_next_frame(data, blit=False) fname = '%s%04d.png' % (frame_prefix, idx) fnames.append(fname) verbose.report('Animation.save: saved frame %d to fname=%s'%(idx, fname), level='debug') self._fig.savefig(fname) self._make_movie(filename, fps, codec, frame_prefix) #Delete temporary files if clear_temp: import os verbose.report('Animation.save: clearing temporary fnames=%s'%str(fnames), level='debug') for fname in fnames: os.remove(fname) # Reconnect signal for first draw if necessary if reconnect_first_draw: self._first_draw_id = self._fig.canvas.mpl_connect('draw_event', self._start) def ffmpeg_cmd(self, fname, fps, codec, frame_prefix): # Returns the command line parameters for subprocess to use # ffmpeg to create a movie return ['ffmpeg', '-y', '-r', str(fps), '-b', '1800k', '-i', '%s%%04d.png' % frame_prefix, fname] def mencoder_cmd(self, fname, fps, codec, frame_prefix): # Returns the command line parameters for subprocess to use # mencoder to create a movie return ['mencoder', 'mf://%s*.png' % frame_prefix, '-mf', 'type=png:fps=%d' % fps, '-ovc', 'lavc', '-lavcopts', 'vcodec=%s' % codec, '-oac', 'copy', '-o', fname] def _make_movie(self, fname, fps, codec, frame_prefix, cmd_gen=None): # Uses subprocess to call the program for assembling frames into a # movie file. *cmd_gen* is a callable that generates the sequence # of command line arguments from a few configuration options. from subprocess import Popen, PIPE if cmd_gen is None: cmd_gen = self.ffmpeg_cmd command = cmd_gen(fname, fps, codec, frame_prefix) verbose.report('Animation._make_movie running command: %s'%' '.join(command)) proc = Popen(command, shell=False, stdout=PIPE, stderr=PIPE) proc.wait() def _step(self, *args): ''' Handler for getting events. By default, gets the next frame in the sequence and hands the data off to be drawn. ''' # Returns True to indicate that the event source should continue to # call _step, until the frame sequence reaches the end of iteration, # at which point False will be returned. try: framedata = self.frame_seq.next() self._draw_next_frame(framedata, self._blit) return True except StopIteration: return False def new_frame_seq(self): 'Creates a new sequence of frame information.' # Default implementation is just an iterator over self._framedata return iter(self._framedata) def new_saved_frame_seq(self): 'Creates a new sequence of saved/cached frame information.' # Default is the same as the regular frame sequence return self.new_frame_seq() def _draw_next_frame(self, framedata, blit): # Breaks down the drawing of the next frame into steps of pre- and # post- draw, as well as the drawing of the frame itself. self._pre_draw(framedata, blit) self._draw_frame(framedata) self._post_draw(framedata, blit) def _init_draw(self): # Initial draw to clear the frame. Also used by the blitting code # when a clean base is required. pass def _pre_draw(self, framedata, blit): # Perform any cleaning or whatnot before the drawing of the frame. # This default implementation allows blit to clear the frame. if blit: self._blit_clear(self._drawn_artists, self._blit_cache) def _draw_frame(self, framedata): # Performs actual drawing of the frame. raise NotImplementedError('Needs to be implemented by subclasses to' ' actually make an animation.') def _post_draw(self, framedata, blit): # After the frame is rendered, this handles the actual flushing of # the draw, which can be a direct draw_idle() or make use of the # blitting. if blit and self._drawn_artists: self._blit_draw(self._drawn_artists, self._blit_cache) else: self._fig.canvas.draw_idle() # The rest of the code in this class is to facilitate easy blitting def _blit_draw(self, artists, bg_cache): # Handles blitted drawing, which renders only the artists given instead # of the entire figure. updated_ax = [] for a in artists: # If we haven't cached the background for this axes object, do # so now. This might not always be reliable, but it's an attempt # to automate the process. if a.axes not in bg_cache: bg_cache[a.axes] = a.figure.canvas.copy_from_bbox(a.axes.bbox) a.axes.draw_artist(a) updated_ax.append(a.axes) # After rendering all the needed artists, blit each axes individually. for ax in set(updated_ax): ax.figure.canvas.blit(ax.bbox) def _blit_clear(self, artists, bg_cache): # Get a list of the axes that need clearing from the artists that # have been drawn. Grab the appropriate saved background from the # cache and restore. axes = set(a.axes for a in artists) for a in axes: a.figure.canvas.restore_region(bg_cache[a]) def _setup_blit(self): # Setting up the blit requires: a cache of the background for the # axes self._blit_cache = dict() self._drawn_artists = [] self._resize_id = self._fig.canvas.mpl_connect('resize_event', self._handle_resize) self._post_draw(None, self._blit) def _handle_resize(self, *args): # On resize, we need to disable the resize event handling so we don't # get too many events. Also stop the animation events, so that # we're paused. Reset the cache and re-init. Set up an event handler # to catch once the draw has actually taken place. self._fig.canvas.mpl_disconnect(self._resize_id) self.event_source.stop() self._blit_cache.clear() self._init_draw() self._resize_id = self._fig.canvas.mpl_connect('draw_event', self._end_redraw) def _end_redraw(self, evt): # Now that the redraw has happened, do the post draw flushing and # blit handling. Then re-enable all of the original events. self._post_draw(None, self._blit) self.event_source.start() self._fig.canvas.mpl_disconnect(self._resize_id) self._resize_id = self._fig.canvas.mpl_connect('resize_event', self._handle_resize) class TimedAnimation(Animation): ''' :class:`Animation` subclass that supports time-based animation, drawing a new frame every *interval* milliseconds. *repeat* controls whether the animation should repeat when the sequence of frames is completed. *repeat_delay* optionally adds a delay in milliseconds before repeating the animation. ''' def __init__(self, fig, interval=200, repeat_delay=None, repeat=True, event_source=None, *args, **kwargs): # Store the timing information self._interval = interval self._repeat_delay = repeat_delay self.repeat = repeat # If we're not given an event source, create a new timer. This permits # sharing timers between animation objects for syncing animations. if event_source is None: event_source = fig.canvas.new_timer() event_source.interval = self._interval Animation.__init__(self, fig, event_source=event_source, *args, **kwargs) def _step(self, *args): ''' Handler for getting events. ''' # Extends the _step() method for the Animation class. If # Animation._step signals that it reached the end and we want to repeat, # we refresh the frame sequence and return True. If _repeat_delay is # set, change the event_source's interval to our loop delay and set the # callback to one which will then set the interval back. still_going = Animation._step(self, *args) if not still_going and self.repeat: if self._repeat_delay: self.event_source.remove_callback(self._step) self.event_source.add_callback(self._loop_delay) self.event_source.interval = self._repeat_delay self.frame_seq = self.new_frame_seq() return True else: return still_going def _stop(self, *args): # If we stop in the middle of a loop delay (which is relatively likely # given the potential pause here, remove the loop_delay callback as # well. self.event_source.remove_callback(self._loop_delay) Animation._stop(self) def _loop_delay(self, *args): # Reset the interval and change callbacks after the delay. self.event_source.remove_callback(self._loop_delay) self.event_source.interval = self._interval self.event_source.add_callback(self._step) class ArtistAnimation(TimedAnimation): ''' Before calling this function, all plotting should have taken place and the relevant artists saved. frame_info is a list, with each list entry a collection of artists that represent what needs to be enabled on each frame. These will be disabled for other frames. ''' def __init__(self, fig, artists, *args, **kwargs): # Internal list of artists drawn in the most recent frame. self._drawn_artists = [] # Use the list of artists as the framedata, which will be iterated # over by the machinery. self._framedata = artists TimedAnimation.__init__(self, fig, *args, **kwargs) def _init_draw(self): # Make all the artists involved in *any* frame invisible axes = [] for f in self.new_frame_seq(): for artist in f: artist.set_visible(False) # Assemble a list of unique axes that need flushing if artist.axes not in axes: axes.append(artist.axes) # Flush the needed axes for ax in axes: ax.figure.canvas.draw() def _pre_draw(self, framedata, blit): ''' Clears artists from the last frame. ''' if blit: # Let blit handle clearing self._blit_clear(self._drawn_artists, self._blit_cache) else: # Otherwise, make all the artists from the previous frame invisible for artist in self._drawn_artists: artist.set_visible(False) def _draw_frame(self, artists): # Save the artists that were passed in as framedata for the other # steps (esp. blitting) to use. self._drawn_artists = artists # Make all the artists from the current frame visible for artist in artists: artist.set_visible(True) class FuncAnimation(TimedAnimation): ''' Makes an animation by repeatedly calling a function *func*, passing in (optional) arguments in *fargs*. *frames* can be a generator, an iterable, or a number of frames. *init_func* is a function used to draw a clear frame. If not given, the results of drawing from the first item in the frames sequence will be used. ''' def __init__(self, fig, func, frames=None ,init_func=None, fargs=None, save_count=None, **kwargs): if fargs: self._args = fargs else: self._args = () self._func = func # Amount of framedata to keep around for saving movies. This is only # used if we don't know how many frames there will be: in the case # of no generator or in the case of a callable. self.save_count = save_count # Set up a function that creates a new iterable when needed. If nothing # is passed in for frames, just use itertools.count, which will just # keep counting from 0. A callable passed in for frames is assumed to # be a generator. An iterable will be used as is, and anything else # will be treated as a number of frames. if frames is None: self._iter_gen = itertools.count elif callable(frames): self._iter_gen = frames elif iterable(frames): self._iter_gen = lambda: iter(frames) self.save_count = len(frames) else: self._iter_gen = lambda: iter(range(frames)) self.save_count = frames # If we're passed in and using the default, set it to 100. if self.save_count is None: self.save_count = 100 self._init_func = init_func # Needs to be initialized so the draw functions work without checking self._save_seq = [] TimedAnimation.__init__(self, fig, **kwargs) # Need to reset the saved seq, since right now it will contain data # for a single frame from init, which is not what we want. self._save_seq = [] def new_frame_seq(self): # Use the generating function to generate a new frame sequence return self._iter_gen() def new_saved_frame_seq(self): # Generate an iterator for the sequence of saved data. If there are # no saved frames, generate a new frame sequence and take the first # save_count entries in it. if self._save_seq: return iter(self._save_seq) else: return itertools.islice(self.new_frame_seq(), self.save_count) def _init_draw(self): # Initialize the drawing either using the given init_func or by # calling the draw function with the first item of the frame sequence. # For blitting, the init_func should return a sequence of modified # artists. if self._init_func is None: self._draw_frame(self.new_frame_seq().next()) else: self._drawn_artists = self._init_func() def _draw_frame(self, framedata): # Save the data for potential saving of movies. self._save_seq.append(framedata) # Make sure to respect save_count (keep only the last save_count around) self._save_seq = self._save_seq[-self.save_count:] # Call the func with framedata and args. If blitting is desired, # func needs to return a sequence of any artists that were modified. self._drawn_artists = self._func(framedata, *self._args)