A function to animate cubes

animate_cube.py

'''
Created on Jun 6, 2013

Function to allow the easy animation of cubes

@author: nrobin
'''

import os
import shutil

import numpy as np
import matplotlib.pyplot as plt

import iris


def animate_cube(cube, plot_func, disp_dims, save_dir, save_name, blend=True, ax_hook=None, window_months=None, *args, **kwargs):
    """
    This function allows cubes to be easily animated. The overhead for blending/animating
    is high, so it should normally be run on a server machine.
    
    Args:
    
    * cube:
        The cube to animate
        
    * plot_func:
        A function to use to render the plots e.g. iris.quickplot.pcolormesh
        
    * disp_dims:
        A list of dimensions to plot over. This should be compatible with
        the plot_func. There should be one omitted dimension which is animated.
        
    * save_dir:
        String directory to save the plot to
        
    * save_name:
        The name of the video (without file ending)
        
    * blend:
        A boolean flag to indicate weather to blend
        
    * ax_hook:
        A function that takes a matplotlib.axis object as its only argument.
        Allows formatting to be performed on the plots
        
    * window_months:
        Integer number of months to smooth the data over
        
    *args and **kwargs are passed to the plot function
    
    """
    
    temp_dir = save_dir+'temp/'
    blended_dir = temp_dir+'blended/'
    try:
        os.mkdir(temp_dir)
        os.mkdir(blended_dir)
    except OSError:
        shutil.rmtree(temp_dir)
        os.mkdir(temp_dir)
        os.mkdir(blended_dir)
        
    print "Making plots"
    
    if window_months:
        cube.rolling_window('time', iris.analysis.MEAN, window_months)
        
    try:
        iris.coord_categorisation.add_month(cube, 'time')
    except BaseException as e:
        print str(e)
    
    if 'vmin' not in kwargs.keys():
        kwargs['vmin'] = np.min(cube.data)
    if 'vmax' not in kwargs.keys():
        kwargs['vmax'] = np.max(cube.data)
         
    for i, frameSlice in enumerate(cube.slices(disp_dims)):
        ax = plot_func(frameSlice, **kwargs)
        axes_hook = kwargs.pop('axes_hook', None)
        if axes_hook is not None:
            axes_hook(ax)
        plt.savefig(temp_dir+'plot_%0.3i' % i)
        plt.close()
    
    if blend:
        print "Blending plots"
        os.system("convert {0}*.png -delay 10 -morph 5 {1}%05d.png".format(temp_dir, blended_dir))
    
    print "Animating plots"
    os.system("ffmpeg -r 10 -qscale 2  -i {0}%05d.png -y {1}{2}.mp4".format(blended_dir, save_dir, save_name))
    
    
if __name__=='__main__':
    import iris.quickplot as qplt
    
    def ax_hook(ax):
        ax.set_title("test")
    
    cube = iris.load_cube(iris.sample_data_path('GloSea4', 'ensemble_001.pp'))
    animate_cube(cube, qplt.pcolormesh, ['latitude', 'longitude'], './', 'test', ax_hook=ax_hook, window_months=3)

Following on from our look round the standard library last week, I suggest you use the subprocess module rather than os.system and investigate the tempfilemodule.