How to add a string as the artist in matplotlib legend?

I am trying to create a legend in a python figure where the artist is a string (a single letter) which is then labelled. For example I would like a legend for the following figure:

import numpy as np
import matplotlib.pyplot as plt
import string

N = 7
x = np.random.rand(N)
y = np.random.rand(N)
colors = np.random.rand(N)
area = np.pi * (15 * np.random.rand(N))**2 

plt.scatter(x, y, s=area, c=colors, alpha=0.5)
for i,j in enumerate(zip(x,y)):
    plt.annotate(list(string.ascii_uppercase)[i],xy=j)
plt.show()

Where the legend is something like:

A - Model Name A

B - Model Name B

C - Model Name C

D - Model Name D

etc.etc.

What I can't work out how to do is place 'A', 'B', .... as the artist for the legend text. I can see how you would use a line or Patch, or something similar. But in general is there a way to use a string as the artist instead of, say, a line?

Solution 1:

I don't think there's a legend handler for text (see the list of available ones here). But you can implement your own custom legend handler. Here I'll just modify the example at the above link:

import matplotlib.pyplot as plt
import matplotlib.text as mpl_text

class AnyObject(object):
    def __init__(self, text, color):
        self.my_text = text
        self.my_color = color

class AnyObjectHandler(object):
    def legend_artist(self, legend, orig_handle, fontsize, handlebox):
        print orig_handle
        x0, y0 = handlebox.xdescent, handlebox.ydescent
        width, height = handlebox.width, handlebox.height
        patch = mpl_text.Text(x=0, y=0, text=orig_handle.my_text, color=orig_handle.my_color, verticalalignment=u'baseline', 
                                horizontalalignment=u'left', multialignment=None, 
                                fontproperties=None, rotation=45, linespacing=None, 
                                rotation_mode=None)
        handlebox.add_artist(patch)
        return patch

obj_0 = AnyObject("A", "purple")
obj_1 = AnyObject("B", "green")

plt.legend([obj_0, obj_1], ['Model Name A', 'Model Name B'],
           handler_map={obj_0:AnyObjectHandler(), obj_1:AnyObjectHandler()})

plt.show()

Solution 2:

The solution will depend on whether you already have texts in the axes that should appear in the legend as well, or whether those are independent on anything you have in the axes.

A. Existing texts or annotation

If you already have texts or annotations in the axes, you can provide them as handles to the legend. A new TextHandlerA that is registered to the Legend class takes those Texts as input. The respective label is taken from the artist as usual, via the label argument.

import numpy as np
import matplotlib.pyplot as plt
import string

from matplotlib.legend_handler import HandlerBase
from matplotlib.text import Text, Annotation
from matplotlib.legend import Legend


class TextHandlerA(HandlerBase):
    def create_artists(self, legend, artist ,xdescent, ydescent,
                        width, height, fontsize, trans):
        tx = Text(width/2.,height/2, artist.get_text(), fontsize=fontsize,
                  ha="center", va="center", fontweight="bold")
        return [tx]

Legend.update_default_handler_map({Text : TextHandlerA()})

N = 7
x = np.random.rand(N)*.7
y = np.random.rand(N)*.7
colors = np.random.rand(N)

handles = list(string.ascii_uppercase) 
labels = [f"Model Name {c}" for c in handles]

fig, ax = plt.subplots()
ax.scatter(x, y, s=100, c=colors, alpha=0.5)
for i, xy in enumerate(zip(x, y)):
    ax.annotate(handles[i], xy=xy, label= labels[i])

ax.legend(handles=ax.texts)
plt.show()

B. Legend from list of strings.

If you want legend entries that are not themselves texts in the axes, you can create them from a list of strings. In this case the TextHandlerB takes the string as input. In that case the legend needs to be called with two lists of strings, one for the handles, and one for the labels.

import numpy as np
import matplotlib.pyplot as plt
import string

from matplotlib.legend_handler import HandlerBase
from matplotlib.text import Text
from matplotlib.legend import Legend


class TextHandlerB(HandlerBase):
    def create_artists(self, legend, text ,xdescent, ydescent,
                        width, height, fontsize, trans):
        tx = Text(width/2.,height/2, text, fontsize=fontsize,
                  ha="center", va="center", fontweight="bold")
        return [tx]

Legend.update_default_handler_map({str : TextHandlerB()})

N = 7
x = np.random.rand(N)*.7
y = np.random.rand(N)*.7
colors = np.random.rand(N)

handles = list(string.ascii_uppercase)[:N] 
labels = [f"Model Name {c}" for c in handles]

fig, ax = plt.subplots()
ax.scatter(x, y, s=100, c=colors, alpha=0.5)
for i, xy in enumerate(zip(x, y)):
    ax.annotate(handles[i], xy=xy)

ax.legend(handles=handles, labels=labels)
plt.show()

In both cases the output is