Create separate function add_colorbar()

Author: mgeier

examples/soundfigures.py

@@ -44,7 +44,7 @@
 
 # plot and save level of synthesized sound field
 plt.figure(figsize=(12.5, 12.5))
-im = sfs.plot.level(p, grid, xnorm=[0, -2.2, 0], colorbar_label='dB',
-                    vmin=-50, vmax=0)
+im = sfs.plot.level(p, grid, xnorm=[0, -2.2, 0], vmin=-50, vmax=0,
+                    colorbar_kwargs=dict(label='dB'))
 plt.title('Level of Synthesized Sound Field')
 plt.savefig('soundfigure_level.png')

sfs/plot.py

@@ -4,7 +4,7 @@
 from matplotlib.patches import PathPatch
 from matplotlib.path import Path
 from matplotlib.collections import PatchCollection
-from mpl_toolkits.axes_grid1 import make_axes_locatable, axes_size
+from mpl_toolkits import axes_grid1
 from mpl_toolkits.mplot3d import Axes3D
 import numpy as np
 from . import util
@@ -166,10 +166,9 @@ def loudspeaker_3d(x0, n0, a0=None, w=0.08, h=0.08):
     fig.show()
 
 
-def soundfield(p, grid, xnorm=None, cmap='coolwarm_clip', colorbar=True,
-               colorbar_aspect=20, colorbar_pad=0.5, colorbar_label='',
-               xlabel=None, ylabel=None, ax=None, vmin=-2.0, vmax=2.0,
-               **kwargs):
+def soundfield(p, grid, xnorm=None, cmap='coolwarm_clip', vmin=-2.0, vmax=2.0,
+               xlabel=None, ylabel=None, colorbar=True, colorbar_kwargs={},
+               ax=None, **kwargs):
     """Two-dimensional plot of sound field.
 
     Parameters
@@ -219,16 +218,8 @@ def soundfield(p, grid, xnorm=None, cmap='coolwarm_clip', colorbar=True,
         :func:`matplotlib.pyplot.ylabel`.
     colorbar : bool, optional
         If ``False``, no colorbar is created.
-    colorbar_aspect : float, optional
-        Aspect ratio of the colorbar, see
-        :func:`matplotlib.pyplot.colorbar`.
-        Strictly speaking, since the colorbar is vertical, it's actually
-        the inverse of the aspect ratio.
-    colorbar_pad : float, optional
-        Space between image plot and colorbar, as a fraction of the
-        width of the colorbar.
-    colorbar_label : str, optional
-        Label of the colorbar.
+    colorbar_kwargs : dict, optional
+        Further colorbar arguments, see :func:`add_colorbar`.
     ax : Axes, optional
         If given, the plot is created on `ax` instead of the current
         axis (see :func:`matplotlib.pyplot.gca`).
@@ -294,12 +285,7 @@ def soundfield(p, grid, xnorm=None, cmap='coolwarm_clip', colorbar=True,
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
     if colorbar:
-        divider = make_axes_locatable(ax)
-        width = axes_size.AxesY(ax, aspect=1/colorbar_aspect)
-        pad = axes_size.Fraction(colorbar_pad, width)
-        cax = divider.append_axes("right", size=width, pad=pad)
-        ax.figure.colorbar(im, cax=cax, label=colorbar_label)
-        plt.sca(ax)
+        add_colorbar(im, **colorbar_kwargs)
     return im
 
 
@@ -344,3 +330,43 @@ def particles(x, trim=None, ax=None, xlabel='x (m)', ylabel='y (m)',
         ax.set_xlabel(xlabel)
     if ylabel:
         ax.set_ylabel(ylabel)
+
+
+def add_colorbar(im, aspect=20, pad=0.5, **kwargs):
+    """Add a vertical color bar to a plot.
+
+    Parameters
+    ----------
+    im : ScalarMappable
+        The output of :func:`sfs.plot.soundfield`,
+        :func:`sfs.plot.level` or any other
+        :class:`matplotlib.cm.ScalarMappable`.
+    aspect : float, optional
+        Aspect ratio of the colorbar.  Strictly speaking, since the
+        colorbar is vertical, it's actually the inverse of the aspect
+        ratio.
+    pad : float, optional
+        Space between image plot and colorbar, as a fraction of the
+        width of the colorbar.
+
+        .. note:: The `pad` argument of
+                  :meth:`matplotlib.figure.Figure.colorbar` has a
+                  slightly different meaning ("fraction of original
+                  axes")!
+    \**kwargs
+        All further arguments are forwarded to
+        :meth:`matplotlib.figure.Figure.colorbar`.
+
+    See Also
+    --------
+    matplotlib.pyplot.colorbar
+
+    """
+    ax = im.axes
+    divider = axes_grid1.make_axes_locatable(ax)
+    width = axes_grid1.axes_size.AxesY(ax, aspect=1/aspect)
+    pad = axes_grid1.axes_size.Fraction(pad, width)
+    current_ax = plt.gca()
+    cax = divider.append_axes("right", size=width, pad=pad)
+    plt.sca(current_ax)
+    return ax.figure.colorbar(im, cax=cax, orientation='vertical', **kwargs)