Display matplotlib legend element as 2D line of colormap

Question

I wish to modify the 2D line in my legend to plot as line segments (or another method like patches) that will display the range of my colormap (here viridis_r) instead of a singular color. While the third variable (radius) is included in the colorbar, having it displayed in the legend as well will be informative when I add more complications to the plot. Thanks!

fig, ax = plt.subplots() 

radii = [1,2,3,4,5]
angle = np.linspace(0, 2 * np.pi, 150)  

cmap = plt.get_cmap('viridis_r')
norm = plt.Normalize(radii[0], radii[-1])    
m = plt.cm.ScalarMappable(cmap=cmap)
m.set_array(radii)

for radius in radii: 
    x = radius * np.cos(angle) 
    y = radius * np.sin(angle)  
    ax.plot(x, y, color=cmap(norm(radius))) 

radius_2Dline = plt.Line2D((0, 1), (0, 0), color='k', linewidth=2)
ax.legend([radius_2Dline],['Radius'], loc='best')

ax.set_aspect( 1 ) 
fig.colorbar(m).set_label('Radius', size=15) 
plt.show() 

Answer 1

The following approach uses the "tuple legend handler". That handler puts a list of legend handles (in this case the circles drawn via ax.plot). Setting ndivide=None will draw one short line for each element in the list. The padding can be set to 0 to avoid gaps between these short lines. The default handlelength might be too small to properly see these special handles; therefore, the example code below increases it a bit.

import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerTuple
import numpy as np

fig, ax = plt.subplots()

radii = [1, 2, 3, 4, 5]
angle = np.linspace(0, 2 * np.pi, 150)

cmap = plt.get_cmap('viridis_r')
norm = plt.Normalize(radii[0], radii[-1])

lines = []  # list of lines to be used for the legend
for radius in radii:
     x = radius * np.cos(angle)
     y = radius * np.sin(angle)
     line, = ax.plot(x, y, color=cmap(norm(radius)))
     lines.append(line)

ax.legend(handles=[tuple(lines)], labels=['Radius'],
          handlelength=3, handler_map={tuple: HandlerTuple(ndivide=None, pad=0)})
ax.set_aspect('equal')
plt.tight_layout()
plt.show()