Merge pull request #159 from ubermag/mpl-cbar

Scale colorbar to axes height - mpl

discretisedfield/__init__.py

@@ -5,6 +5,7 @@
 import pkg_resources
 import pytest
 
+from . import tools
 from .field import Field
 from .field_rotator import FieldRotator
 from .interact import interact

discretisedfield/plotting/mpl_field.py

@@ -4,6 +4,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import ubermagutil.units as uu
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 import discretisedfield as df
 import discretisedfield.util as dfu
@@ -302,9 +303,7 @@ def scalar(
         cp = ax.imshow(np.transpose(values), origin="lower", extent=extent, **kwargs)
 
         if colorbar:
-            cbar = plt.colorbar(cp, ax=ax)
-            if colorbar_label is not None:
-                cbar.ax.set_ylabel(colorbar_label)
+            self._add_colorbar(ax, cp, colorbar_label)
 
         self._axis_labels(ax, multiplier)
 
@@ -662,9 +661,7 @@ def vector(
 
         ax.set_aspect("equal")
         if colorbar and use_color:
-            cbar = plt.colorbar(cp, ax=ax)
-            if colorbar_label is not None:
-                cbar.ax.set_ylabel(colorbar_label)
+            self._add_colorbar(ax, cp, colorbar_label)
 
         self._axis_labels(ax, multiplier)
 
@@ -797,9 +794,7 @@ def contour(
         ax.set_aspect("equal")
 
         if colorbar:
-            cbar = plt.colorbar(cp, ax=ax)
-            if colorbar_label is not None:
-                cbar.ax.set_ylabel(colorbar_label)
+            self._add_colorbar(ax, cp, colorbar_label)
 
         self._axis_labels(ax, multiplier)
 
@@ -849,3 +844,10 @@ def __dir__(self):
             dirlist.remove(attr)
 
         return dirlist
+
+    def _add_colorbar(self, ax, cp, colorbar_label):
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        cbar = plt.colorbar(cp, cax=cax)
+        if colorbar_label is not None:
+            cbar.ax.set_ylabel(colorbar_label)