Merge branch 'ColorBarFix' into 'develop'

Color bar fixed to align with the height of the heatmaps

See merge request !374

Wrap/Python/ba_plot.py

@@ -18,6 +18,9 @@ try:  # workaround for build servers
     import numpy as np
     import matplotlib as mpl
     from matplotlib import pyplot as plt
+    from matplotlib import gridspec, colors
+    from mpl_toolkits.axes_grid1 import make_axes_locatable, axes_size
+
 except Exception as e:
     print("In ba_plot.py: {:s}".format(str(e)))
 
@@ -108,22 +111,29 @@ def plot_array(array, axes_limits=None, **kwargs):
     cmap = kwargs.pop('cmap', CMAP)
     withCBar = kwargs.pop('with_cb', True)
 
-    im = plt.imshow(array,
+    ax = plt.gca()
+    im = ax.imshow(array,
                     cmap=cmap,
                     norm=norm,
                     aspect=aspect,
                     extent=axes_limits,
                     **kwargs)
-    if withCBar:
-        cb = plt.colorbar(im, pad=0.025)
 
     if xlabel:
         plt.xlabel(xlabel, fontsize=label_fontsize)
     if ylabel:
         plt.ylabel(ylabel, fontsize=label_fontsize)
+    if withCBar:
+        aspect = 20
+        pad_fraction = 3 
+
+        divider = make_axes_locatable(ax)
+        width = axes_size.AxesY(ax, aspect=1./aspect)
+        pad = axes_size.Fraction(pad_fraction, width)
+        cax = divider.append_axes("right", size=width, pad=pad)
+        cb = plt.colorbar(im, cax=cax)
     if zlabel:
         cb.set_label(zlabel, size=label_fontsize)
-
     if title:
         plt.title(title)