Matplotlib: How to plot a small rectangle filled with a colormap as a legend

Question

Instead of a plotting a colorbar next to my plot, I want to plot a small rectangle filled with the colormap as a legend.

I already can plot a small rectangle filled with any color, by doing the following trick:

axis0, = myax.plot([], linewidth=10, color='r')

axis =[axis0]
legend=['mytext']

plt.legend(axis,
           legend)

Can I do the same with a colormap? Thanks!

Answer 1

To my knowledge, there is no way of doing this other than by creating the rectangle and legend from scratch. Here is one approach (based mainly on this answer):

import numpy as np                                    # v 1.19.2
import matplotlib.pyplot as plt                       # v 3.3.2
import matplotlib.patches as patches
from matplotlib.legend_handler import HandlerTuple

rng = np.random.default_rng(seed=1)

ncmaps = 5     # number of colormaps to draw for illustration
ncolors = 100  # number high enough to draw a smooth gradient for each colormap

# Create random list of colormaps and extract list of colors to 
# draw the gradient of each colormap
cmaps_names = list(rng.choice(plt.colormaps(), size=ncmaps))
cmaps = [plt.cm.get_cmap(name) for name in cmaps_names]
cmaps_gradients = [cmap(np.linspace(0, 1, ncolors)) for cmap in cmaps]
cmaps_dict = dict(zip(cmaps_names, cmaps_gradients))

# Create a list of lists of patches representing the gradient of each colormap
patches_cmaps_gradients = []
for cmap_name, cmap_colors in cmaps_dict.items():
    cmap_gradient = [patches.Patch(facecolor=c, edgecolor=c, label=cmap_name)
                     for c in cmap_colors]
    patches_cmaps_gradients.append(cmap_gradient)

# Create custom legend (with a large fontsize to better illustrate the result)
plt.legend(handles=patches_cmaps_gradients, labels=cmaps_names, fontsize=20,
           handler_map={list: HandlerTuple(ndivide=None, pad=0)})

plt.show()

If you plan on doing this for multiple plots, you may want to create a custom legend handler as shown in this answer. You may also want to consider other ways of displaying the colorbar, such as in the examples shown here and here and here.

Documentation: legend guide