Custom categorized legend matplotlib

Question

I made a mock-up version of my dataset to illustrate my problem.

I have a graph like this:

import pandas as pd
import matplotlib.pyplot as plt

data = {'x': [0, 1, 2, 3],
        'y': [3, 2, 1, 0],
        'cat1': ['A', 'B', 'A', 'B'],
        'cat2': ['f', 'g', 'h', 'i']}

df = pd.DataFrame(data)

colors = {'A':'tab:red',
          'B':'tab:blue'}

markers = {'f':"o",
           'g':"v",
           'h':"s",
           'i':"+"}

fig, ax = plt.subplots()
for i in range(df.shape[0]):
  ax.scatter(df.loc[i,'x'], 
             df.loc[i,'y'], 
             color=colors[df.loc[i,'cat1']], 
             marker=markers[df.loc[i,'cat2']],
             label = df.loc[i, 'cat2'])
  
ax.legend() 

But I'm looking for a legend like this:

Could anyone give me some tips on how to solve this? Also, it would be better if the legend in the final plot were in a box outside the plot, on the right side.

Answer 1

To add additional lines (A,B), rearrange the order and move the legend outside the graph, these are the steps you can follow after plotting.

• Create a custom legend entries for the new entries and existing one using Line2D
• Plot the legend and use bbox_to_anchor to move the legend to where you need it. You can adjust the coordinates within bbox if you want to move the position
• Adjust the labels for A, B as these are only text (no marker) so that align horizontally to the middle (you can adjust set_position() if you want to move it further to the left/right

Code

import pandas as pd
import matplotlib.pyplot as plt

data = {'x': [0, 1, 2, 3],
        'y': [3, 2, 1, 0],
        'cat1': ['A', 'B', 'A', 'B'],
        'cat2': ['f', 'g', 'h', 'i']}

df = pd.DataFrame(data)

colors = {'A':'tab:red',
          'B':'tab:blue'}

markers = {'f':"o",
           'h':"s",
           'g':"v",
           'i':"+"}

fig, ax = plt.subplots()
for i in range(df.shape[0]):
    ax.scatter(df.loc[i,'x'], df.loc[i,'y'], 
        color=colors[df.loc[i,'cat1']], 
        marker=markers[df.loc[i,'cat2']],
        label = df.loc[i, 'cat2']

## Create legend handle entries for each of the items
from matplotlib.lines import Line2D
title = Line2D([0], [0], linestyle="none", marker="")
f = Line2D([0], [0], linestyle="none", marker="o", markersize=10, markeredgecolor='tab:red', markerfacecolor="tab:red")
g = Line2D([0], [0], linestyle="none", marker="v", markersize=10,  markeredgecolor='tab:blue', markerfacecolor="tab:blue")
h = Line2D([0], [0], linestyle="none", marker="s", markersize=10,  markeredgecolor='tab:red', markerfacecolor="tab:red")
i = Line2D([0], [0], linestyle="none", marker="+", markersize=10,  markeredgecolor='tab:blue', markerfacecolor="tab:blue")

## Plot in order you want, bbox to set legend box outside
leg=ax.legend((title, f, h, title, g, i), ('A', 'f', 'h', 'B', 'g','i'), bbox_to_anchor=(1.16, 1.03))

## Adjust position of A and B so that they are in middle
for item, label in zip(leg.legendHandles, leg.texts):
    if label._text in ['A', 'B']:
        width=item.get_window_extent(fig.canvas.get_renderer()).width
        label.set_ha('left')
        label.set_position((-width/2,0)) ## Adjust here to move left/right

Plot