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()

Answer 2

I am not sure if this is your goal but here is a stab at it. Following this answer, you can make a 'fake' legend with a colormap.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


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)))

# Set box that will act as a 'fake' legend, 25% width of the 
# x-axis, 15% of y-axis
cbbox = inset_axes(ax,  width="25%", height="15%", loc=2)
cbbox.tick_params(
    axis = 'both',
    left = False,
    top = False,
    right = False,
    bottom = False,
    labelleft = False,
    labeltop = False,
    labelright = False,
    labelbottom = False
)
# Semi-transparent like the usual ax.legend()
cbbox.set_facecolor([1, 1, 1, 0.7])

# Colorbar inside the fake legend box, occupying 85% of the 
# box width and %5 box height
cbaxes = inset_axes(cbbox, width="85%", height="5%", loc=2)
cbar = fig.colorbar(m, cax=cbaxes, orientation='horizontal', 
ticks=[1, 3, 5])
cbar.set_label('Radius', size=9)
cbar.ax.tick_params(labelsize=9)

ax.set_aspect(1)
plt.show()

I was unsuccessful in creating an actual ax.legend() from a LineCollection or a multicolored line - it only plotted one color - so my solution was this 'fake' legend approach. Hope this helps, cheers.