Legend for Lines with Different Endpoint Markers

user_5020803 0
Viewed 200

I need to create a legend for a line segment that has a different marker at the bottom and the top of the line. I am able to create a legend with 1 of the marker symbols repeated but not the two different markers on each end.

Here is a reproducible example.

from matplotlib import pyplot as plt
import numpy as np

#Create some data
x = np.arange(0,11)
y1 = np.sqrt(x/2.)
y2 = x


plt.figure(figsize=(8,8))
ax = plt.subplot(111)

#Plot the lines
for i,x_ in zip(range(11),x):
    ax.plot([x_,x_],[y1[i],y2[i]],c='k')
    
#Plot the end points    
ax.scatter(x,y1,marker='s',c='r',s=100,zorder=10)
ax.scatter(x,y2,marker='o',c='r',s=100,zorder=10)

ax.plot([],[],c='k',marker='o',mfc='r',label='Test Range') #Create a single line for the label
ax.legend(loc=2,numpoints=2,prop={'size':16}) # How can I add a label with different symbols the line segments?

plt.show()

The end product should have a legend with a symbol showing a line connecting a circle and a square.

enter image description here

matplotlib python
1 Answers

I'm afraid you have to combine different patches of mpatches, I'm not sure whether there is a better solution

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as mpatches
from matplotlib.legend_handler import HandlerPatch
from matplotlib.legend_handler import HandlerLine2D


class HandlerCircle(HandlerPatch):
    def create_artists(self,legend,orig_handle,
                       xdescent,ydescent,width,height,fontsize,trans):
        center = 0.5 * width, 0.5 * height
        p = mpatches.Circle(xy=center,radius=width*0.3) 
        self.update_prop(p,orig_handle,legend)
        p.set_transform(trans)
        return [p]
    
class HandlerRectangle(HandlerPatch):
    def create_artists(self, legend, orig_handle,
                       xdescent, ydescent, width, height, fontsize, trans):
        center = 0,height/2-width*0.5/2
        width,height = width*0.5,width*0.5
        p = mpatches.Rectangle(xy=center,width=width,height=width) 
        self.update_prop(p, orig_handle, legend)
        p.set_transform(trans)
        return [p]

fig,ax = plt.subplots(figsize=(12,8))
texts = ['','','Test Range']
line, = ax.plot([],[],c='k')
c = [mpatches.Circle((0.,0.,),facecolor='r',linewidth=.5),
     line,
     mpatches.Rectangle((0.,0.),5,5,facecolor='r',linewidth=.5)]

ax.legend(c,texts,bbox_to_anchor=(.25,.95),loc='center',ncol=3,prop={'size':20},
          columnspacing=-1,handletextpad=.6,
          handler_map={mpatches.Circle: HandlerCircle(), 
                       line: HandlerLine2D(numpoints=0,),
                       mpatches.Rectangle: HandlerRectangle()}).get_frame().set_facecolor('w')

plt.show()

running this script, you will get

output

If you use a different figure size or a different legend size, the settings in my script above may not be optimal. In that case, you can adjust the following parameters:

  1. The centers and the sizes of Circle and Rectangle
  2. columnspacing and handletextpad in ax.legend(...)
user_83668050
Related