Fixes #74 (improving add_colorbar function)

finesse/plotting/tools.py

-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