#
# colourbarbitmap.py - A function which renders a colour bar using
# matplotlib as an RGBA bitmap.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides a single function, :func:`colourBarBitmap`, which uses
:mod:`matplotlib` to plot a colour bar. The colour bar is rendered off-screen
and returned as an RGBA bitmap.
"""
[docs]def colourBarBitmap(cmap,
width,
height,
cmapResolution=256,
negCmap=None,
invert=False,
gamma=1,
ticks=None,
ticklabels=None,
tickalign=None,
label=None,
orientation='vertical',
labelside='top',
alpha=1.0,
fontsize=10,
bgColour=None,
textColour='#ffffff',
scale=1.0):
"""Plots a colour bar using :mod:`matplotlib`.
The rendered colour bar is returned as a RGBA bitmap within a
``numpy.uint8`` array of size :math:`w \\times h \\times 4`, with the
top-left pixel located at index ``[0, 0, :]``.
A rendered colour bar will look something like this:
.. image:: images/colourbarbitmap.png
:scale: 50%
:align: center
:arg cmap: Name of a registered :mod:`matplotlib` colour map.
:arg width: Colour bar width in pixels.
:arg height: Colour bar height in pixels.
:arg cmapResolution: Colour map resolution (number of distinct colours).
:arg negCmap: If provided, two colour maps are drawn, centered at 0.
:arg invert: If ``True``, the colour map is inverted.
:arg gamma: Gamma correction factor - exponentially weights the
colour map scale towards one end.
:arg ticks: Locations of tick labels. Ignored if
``ticklabels is None``.
:arg ticklabels: Tick labels.
:arg tickalign: Tick alignment (one for each tick, either ``'left'``,
``'right'``, or ``'center'``).
:arg label: Text label to show next to the colour bar.
:arg orientation: Either ``vertical`` or ``horizontal``.
:arg labelside: Side of the colour bar to put the label - ``top``,
``bottom``, ``left`` or ``right``. If
``orientation='vertical'``, then ``top``/``bottom``
are interpreted as ``left``/``right`` (and vice-versa
when ``orientation='horizontal'``).
:arg alpha: Colour bar transparency, in the range ``[0.0 - 1.0]``.
:arg fontsize: Label font size in points.
:arg bgColour: Background colour - can be any colour specification
that is accepted by :mod:`matplotlib`.
:arg textColour: Label colour - can be any colour specification that
is accepted by :mod:`matplotlib`.
:arg scale: DPI scaling factor.
"""
# These imports are expensive, so we're
# importing at the function level.
import numpy as np
import matplotlib.backends.backend_agg as mplagg
import matplotlib.figure as mplfig
if orientation not in ['vertical', 'horizontal']:
raise ValueError('orientation must be vertical or '
'horizontal ({})'.format(orientation))
if orientation == 'horizontal':
if labelside == 'left': labelside = 'top'
elif labelside == 'right': labelside = 'bottom'
else:
if labelside == 'top': labelside = 'left'
elif labelside == 'bottom': labelside = 'right'
if labelside not in ['top', 'bottom', 'left', 'right']:
raise ValueError('labelside must be top, bottom, left '
'or right ({})'.format(labelside))
# vertical plots are rendered horizontally,
# and then simply rotated at the end
if orientation == 'vertical':
width, height = height, width
if labelside == 'left': labelside = 'top'
else: labelside = 'bottom'
# Default is 96 dpi to an inch
winches = width / 96.0
hinches = height / 96.0
dpi = scale * 96.0
ncols = cmapResolution
data = genColours(cmap, ncols, invert, alpha, gamma)
if negCmap is not None:
ndata = genColours(negCmap, ncols, not invert, alpha, gamma)
data = np.concatenate((ndata, data), axis=1)
ncols *= 2
# force tick positions to
# the left edge of the
# corresponding colour
if ticks is not None:
ticks = [t - 0.5 / ncols for t in ticks]
fig = mplfig.Figure(figsize=(winches, hinches), dpi=dpi)
canvas = mplagg.FigureCanvasAgg(fig)
ax = fig.add_subplot(111)
if bgColour is not None:
fig.patch.set_facecolor(bgColour)
else:
fig.patch.set_alpha(0)
# draw the colour bar
ax.imshow(data,
aspect='auto',
origin='lower',
interpolation='nearest')
ax.set_yticks([])
ax.tick_params(colors=textColour, labelsize=fontsize, length=0)
if labelside == 'top':
ax.xaxis.tick_top()
ax.xaxis.set_label_position('top')
else:
ax.xaxis.tick_bottom()
ax.xaxis.set_label_position('bottom')
if label is not None:
ax.set_xlabel(label, fontsize=fontsize, color=textColour)
label = ax.xaxis.get_label()
if ticks is None or ticklabels is None:
ax.set_xticks([])
else:
ax.set_xticks(np.array(ticks) * ncols)
ax.set_xticklabels(ticklabels)
ticklabels = ax.xaxis.get_ticklabels()
# Resize the colour bar to make
# space for the ticks and labels
wpad = 5 / float(width)
hpad = 5 / float(height)
totalHeight = 1.0
# Figure out text height in
# pixels (with 6 pixel padding)
fontpix = (fontsize / 72.0) * dpi / scale
textHeight = (fontpix + 6) / float(height)
if label is not None: totalHeight -= textHeight
if ticklabels is not None: totalHeight -= textHeight
if labelside == 'top': bottom, top = 0, totalHeight
else: bottom, top = 1 - totalHeight, 1
fig.subplots_adjust(
left=wpad,
right=1.0 - wpad,
bottom=bottom + hpad,
top=top - hpad)
# Divide the vertical height between the
# colour bar, the label, and the tick labels.
if label is not None:
# I don't understand why, but I have
# to set va to the opposite of what
# I would have thought.
if labelside == 'top':
label.set_va('bottom')
label.set_position((0.5, 1.0))
else:
label.set_va('top')
label.set_position((0.5, 0))
# Adjust tick label horizontal alignment. This
# must be done *after* calling tick_top/tick_bottom,
# as I think the label objects get recreated.
if ticklabels is not None and tickalign is not None:
for l, a in zip(ticklabels, tickalign):
l.set_horizontalalignment(a)
ax.set_xlim((-0.5, ncols - 0.5))
canvas.draw()
buf = canvas.tostring_argb()
ncols, nrows = canvas.get_width_height()
bitmap = np.frombuffer(buf, dtype=np.uint8)
bitmap = bitmap.reshape(nrows, ncols, 4).transpose([1, 0, 2])
# the bitmap is in argb order,
# but we want it in rgba
rgb = bitmap[:, :, 1:]
a = bitmap[:, :, 0]
bitmap = np.dstack((rgb, a))
if orientation == 'vertical':
bitmap = np.flipud(bitmap.transpose([1, 0, 2]))
bitmap = np.rot90(bitmap, 2)
return bitmap
[docs]def genColours(cmap, cmapResolution, invert, alpha, gamma=1):
"""Generate an array containing ``cmapResolution`` colours from the given
colour map object/function.
"""
import numpy as np
import matplotlib.cm as cm
ncols = cmapResolution
cmap = cm.get_cmap(cmap)
idxs = np.linspace(0.0, 1.0, ncols)
if gamma != 1:
idxs = idxs ** gamma
if invert:
idxs = idxs[::-1]
cmap = cmap(idxs)
cmap[:, 3] = alpha
return np.dstack((cmap, cmap)).transpose((2, 0, 1))