Commit ab0baa46 authored by Samuel Rowlinson's avatar Samuel Rowlinson

Fixes #74 (improving add_colorbar function)

parent f43a241f
def add_colorbar(im, aspect=30, pad_fraction=0.5, **kwargs):
def add_colorbar(im, **kwargs):
"""Adds a vertical color bar to an image plot.
Parameters
----------
im : :class:`matplotlib.image.AxesImage`
An image plot, i.e. return value of a call to ``plt.imshow()``.
An image plot, e.g. return value of a call to ``plt.imshow()``.
Returns
-------
Handle to the colorbar.
"""
## Credit to: https://joseph-long.com/writing/colorbars/
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.pyplot as plt
from mpl_toolkits import axes_grid1
divider = axes_grid1.make_axes_locatable(im.axes)
width = axes_grid1.axes_size.AxesY(im.axes, aspect=1.0 / aspect)
pad = axes_grid1.axes_size.Fraction(pad_fraction, width)
current_ax = plt.gca()
cax = divider.append_axes("right", size=width, pad=pad)
plt.sca(current_ax)
return im.axes.figure.colorbar(im, cax=cax, **kwargs)
last_axes = plt.gca()
ax = im.axes
fig = ax.figure
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
cbar = fig.colorbar(im, cax=cax, **kwargs)
plt.sca(last_axes)
return cbar
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment