Python Snippets

Package abbreviations

Unless otherwise specified, the snippets use the following conventions for package abbreviations

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

Improving exponential notation in colour bars

The default behaviour for handling exponents in colour bars isn't the prettiest, and can sometimes overlap the plot window itself. While there are options for shifting it exponent around to fit better, I find it looks a lot better to simply place it in the ylabel for the colour bar. This also provides a nice opportunity to add units!

Function Definition

First, we'll define a function, because it keeps things tidy. To use it, simply pass the handle for a colour bar to the function ScientificCbar. It has some option arguments that may be useful.

1. units (string) Simply set this argument to the desired string to have it included in the colour bar. The default is the empty string.
2. orientation (string) Specifies whether or not the colour bar should be horizontal or vertical. Default is 'vertical'.
3. centre (boolean) If True (the default), will forcibly centre the colour bar around 0.

Notes:

1. If you need more significant digits, modify the format string {0:.2g} (lines 28 and 33) by changing the 2 to be the desired number of digits
2. This setting only permits 5 ticks along the colour bar: this is to avoid clutter while giving basic information on scaling. If you need / want more, simply modify line 15.

import matplotlib as mpl
import numpy as np

def ScientificCbar(cbar, 
        units='', 
        orientation='vertical', 
        centre=True):

    # If requested, centre the colour bar
    if centre:
        cv = np.max(np.abs(cbar.mappable.get_clim()))
        cbar.mappable.set_clim(-cv, cv)

    # Limit the number of ticks on the colour bar
    tick_locator = mpl.ticker.MaxNLocator(nbins=5)
    cbar.locator = tick_locator
    cbar.update_ticks()
    ticks = cbar.get_ticks()

    # Re-scale the values to avoid the poorly-placed exponent
    #   above the colour bar
    scale = np.log10(np.max(np.abs(ticks)))
    scale = np.floor(scale)

    # Instead, simply add a ylabel to the colour bar giving the scale.
    if orientation == 'vertical':
        if scale != 0.:
            cbar.ax.set_yticklabels(["{0:.2g}".format(tick/(10**scale)) for tick in ticks])
        cbar.ax.set_ylabel('$\\times10^{' + '{0:d}'.format(int(scale)) + '}$ ' + units,
                rotation = '-90', labelpad=10)
    elif orientation == 'horizontal':
        if scale != 0.:
            cbar.ax.set_xticklabels(["{0:.2g}".format(tick/(10**scale)) for tick in ticks])
        cbar.ax.set_xlabel('$\\times10^{' + '{0:d}'.format(int(scale)) + '}$ ' + units,
                rotation = '0', labelpad=10)

Sample Usage

# Create colour bar in usual fashion
#   the second and third line are useful
#   for pdf outputs
cbar = plt.colorbar(plot_handle, ...)
cbar.solids.set_rasterized(True)
cbar.solids.set_edgecolor("face")

ScientificCbar(cbar)

Merging Images into a Movie

The following function provides a clean wrapper for merging a collection of images (such as pngs) into a movie (such as an mp4).

Notes:

1. This requires that ffmpeg be installed and accessible at command-line.
2. The first argument, frame_filenames, is a string that indicates the files. For example, if the files are img_000.png, img_001.png, ..., then you would pass img_%03d.png
3. The second argument, movie_name, is the desired output filename. This should include the extension. For example, my_mov.mp4
4. The (optional) argument fps specifies the frame rate out the output. Default is 12 frames per second.
5. The first two arguments use relative paths, so you can merge files in one directory and save the video in another.
6. The log files ffmpeg.log and ffmpeg.err contain the outputs from stdout and stderr, respectively

import subprocess

def merge_to_mp4(frame_filenames, movie_name, fps=12):
    f_log = open("ffmpeg.log", "w")
    f_err = open("ffmpeg.err", "w")
    cmd = ['ffmpeg', '-framerate', str(fps), '-i', frame_filenames, '-y', 
            '-q', '1', '-threads', '0', '-pix_fmt', 'yuv420p', movie_name]
    subprocess.call(cmd, stdout=f_log, stderr=f_err)
    f_log.close()
    f_err.close()