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 centre_val.
4. centre_val (double) The centre value about which the colour will be centred, if centre=True. Default is 0.
5. labelpad (double) Extra space to pad between label and colour bar. The larger it is, the further the label from the colour bar.
6. label (string) Identification label to be prepended to the colour bar label. E.g. the name of the field being plotted.

Notes:

1. If you need more significant digits, modify the format string {0:.2g} (line 31) 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, centre_val=0.,
        labelpad = 20, label=''):

    # If requested, centre the colour bar
    if centre:
        cb_vals = cbar.mappable.get_clim()
        cv = np.max(np.abs(np.array([val - centre_val for val in cb_vals])))
        cbar.mappable.set_clim(centre_val-cv, centre_val+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)

    # Label
    cb_label = '$\\times10^{' + '{0:d}'.format(int(scale)) + '}$ ' + units
    if len(label) > 0:
        cb_label = label + '\n' + cb_label

    # Tick labels
    tick_labels = ["{0:.2g}".format(tick/(10**scale)) for tick in ticks]

    # Instead, simply add a ylabel to the colour bar giving the scale.
    if orientation == 'vertical':
        if scale != 0.:
            cbar.ax.set_yticklabels(tick_labels)
        cbar.ax.set_ylabel(cb_label, rotation = '-90', labelpad = labelpad)
    elif orientation == 'horizontal':
        if scale != 0.:
            cbar.ax.set_xticklabels(tick_labels)
        cbar.ax.set_xlabel(cb_label, rotation = '0', labelpad = labelpad)

Sample Usage / Organizing Multiple Colour Bars

This is just a snippet, not a full working example.

One could certainly just throw the three loop segments into one part, but they're kept separate here because in practice, the contents of the first loop is likely to be a lot messier.

fig, axes = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(7,7))

# Create an array to store the pcolormesh objects
qs = np.array(axes.shape, dtype='object')

# Also create an array to store the colour bars
cbars = np.array(axes.shape, dtype='object')

# Suppose arrays x and y already exist.
# Suppose data is 4D, with the first two corresponding to y,x
#      and the last two being some parameter index
for ii in range(axes.shape[0]):
    for jj in range(axes.shape[1]):
        qs[ii,jj] = axes[ii,jj].pcolormesh(x, y, data[:,:,ii,jj])


for ii in range(axes.shape[0]):
    for jj in range(axes.shape[1]):
        cbars[ii,jj] = plt.colorbar(qs[ii,jj], ax=axes[ii,jj])


for ii in range(axes.shape[0]):
    for jj in range(axes.shape[1]):
        ScientificCbar(cbars[ii,jj])

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()

Plotting Multiple Lines, Coloured by a Colour Bar

A common desire, particularly for plotting a sweep across a parameter (such as a time evolution), is to plot multiple line plots, but to have them be coloured sequentially, following a colour map, and to have a colour bar showing the evolution.

The below example illustrates exactly how to do this.

Lines 8-15 are purely to generate a sample data-set: a right-wards propagating Gaussian, the steadily widens.

Lines 17-20 are used to sub-sample the data so that we don't have too many lines.

Line 25 creates a list of the desired lines sets.

The colour map is specified in line 28.

import numpy as np
import matplotlib
from matplotlib.collections import LineCollection
import matplotlib.pyplot as plt

fig, ax = plt.subplots(1, 1)

Nx, Nt = 100, 100

x = np.linspace(-1, 1, Nx)
t = np.linspace( 0, 5, Nt)

X, T = np.meshgrid(x, t)

h = np.exp(-( (X - T/5 + 0.5) / (0.2 + T/10))**2)

## Determine which times we want to show
start = 0                    # start at the beginning
stop  = Nt//2                # only plot first half of simulation
skip  = (stop-start) // 10   # only use 10 lines to plot evolution

plot_inds = np.arange(start, stop, skip)

# Select out the data that we want to use for creating lines
h_lines = [np.column_stack([x, h[ii,:]]) for ii in plot_inds]

# Create a LineCollection object which holds them together
h_segments = LineCollection(h_lines, cmap=matplotlib.cm.viridis)
           
# Set which variable is used with the colour map (determines colours)
h_segments.set_array(t[plot_inds])
           
# Plot the lines         
ax.add_collection(h_segments)

# Need to manually update the bounds
ax.axis('tight')
    
# Add a colour bar showing time
cbar = plt.colorbar(line_segments, ax = ax)