Python Snippets

From Fluids Wiki
Jump to navigation Jump to search


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)