Useful code snippets: Plotting & Visualization

Matplotlib

Display colorbars for subplots

When dealing with a single set of color-mappable data, you can just call plt.colorbar(), but when using subplots, the colorbar needs to get its own axis. Meaning, you'd need to create an extra subplot just for the colorbar, that's of the right size.

Or, you can use the below snippet (source: https://stackoverflow.com/a/39938019/1399901)

from mpl_toolkits.axes_grid1 import make_axes_locatable

fig, ax = plt.subplots(1, 2)

lines = ax[0].imshow(dataset_1)
divider = make_axes_locatable(ax[0])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(lines, label='Dataset 1 Colorbar', cax=cax)

lines = ax[1].imshow(dataset_2)
divider = make_axes_locatable(ax[1])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(lines, label="Dataset 2 Colorbar", cax=cax)

fig.tight_layout()

This creates a subplot w/ 2 columns. In each column, a heatmap is plotted with imshow, and then a new axis is added for the colobar, that's 5% the width of the subplot.