Custom categorized legend matplotlib-CodePudding

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.

CodePudding user response：

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