Color bar fixed to align with the height of the heatmaps

The color bar for some non quadratic heatmaps was not aligned with the plot (exceeding on top and bottom). ba_plot.py was changed to align the color bar independent of the heatmap size.