Search code examples
pythonmatplotlibcolorbar

Prevent stretching of colorbar to width of plot axis in matplotlib


enter image description here

The plot above has been produced using mpl_toolkits and matplotlib.colorbar.ColorbarBase, because I needed to customise the colormap and colorbar for my discrete dataset, as shown below:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colorbar
from matplotlib.collections import LineCollection
from matplotlib.colors import BoundaryNorm, ListedColormap
from mpl_toolkits.axes_grid1 import make_axes_locatable
import random

x = np.arange(1, 1142)
y = np.zeros(len(x))
z = []
for _ in range(len(x)):
    z.append(random.randint(-1, 5))
z = np.array(z)

cmap = ListedColormap(['#FF0000', '#D52B00', '#AA5500', '#808000', '#55AA00', '#2BD500', '#00FF00'])
norm = BoundaryNorm([-1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5], cmap.N)

points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(z)
lc.set_linewidth(10)

plt.gca().add_collection(lc)

plt.xlim(x.min() - (x.max() * 0.05), x.max() + (x.max() * 0.05))
plt.ylim(-1.1, 1.1)

plt.tick_params(axis='both', which='both', bottom=False, labelbottom=False, left=False, labelleft=False)

divider = make_axes_locatable(plt.gca())
ax_cb = divider.append_axes('bottom', size="2%", pad=-0.5)
cb = colorbar.ColorbarBase(ax_cb, cmap=cmap, norm=norm, orientation='horizontal', ticks=[-1, 0, 1, 2, 3, 4, 5])
cb.ax.set_yticklabels(['-1', '0', '1', '2', '3', '4', '5'])
plt.gcf().add_axes(ax_cb)

plt.show()

This solution was based on the example here.

My question, is how can I make the colorbar shorter, so that it doesn't stretch across the entire width of the plot axis?


Solution

  • The problem is that the divider created via make_axes_locatable makes sure the new axes are exactly as large as the one from which it is cut. That is the main aim of this function; but here it's somehow orthorgonal to the desire to have a different size.

    The solution would be to not use this kind of divider and create the colorbar in the usual fashion via plt.colorbar or fig.colorbar. This allows to use the arguments shrink and aspect. Since there is 5% margin on each side of the data, you may want to shrink the colorbar to 90% percent.

    plt.colorbar(sm, orientation='horizontal', pad=-0.2,  shrink=0.9, aspect=30,
                 ticks=[-1, 0, 1, 2, 3, 4, 5])
    

    Complete code:

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.collections import LineCollection
    from matplotlib.colors import BoundaryNorm, ListedColormap
    from matplotlib.cm import ScalarMappable
    
    x = np.arange(1, 1142)
    y = np.zeros(len(x))
    z = np.random.randint(-1, 5, size=x.shape)
    
    cmap = ListedColormap(['#FF0000', '#D52B00', '#AA5500', '#808000', '#55AA00', '#2BD500', '#00FF00'])
    norm = BoundaryNorm([-1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5], cmap.N)
    
    points = np.array([x, y]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    
    lc = LineCollection(segments, cmap=cmap, norm=norm)
    lc.set_array(z)
    lc.set_linewidth(10)
    
    plt.gca().add_collection(lc)
    
    plt.xlim(x.min() - (x.max() * 0.05), x.max() + (x.max() * 0.05))
    plt.ylim(-1.1, 1.1)
    
    plt.tick_params(axis='both', which='both', bottom=False, labelbottom=False, 
                    left=False, labelleft=False)
    
    sm = ScalarMappable(norm=norm, cmap=cmap)
    sm.set_array([])
    
    plt.colorbar(sm, orientation='horizontal', pad=-0.2,  shrink=0.9, aspect=30,
                 ticks=[-1, 0, 1, 2, 3, 4, 5])
    
    plt.show()
    

    enter image description here