如何在Seaborn facetgrid条形图上添加图例

Question

I have the following code:

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
matplotlib.style.use('ggplot')
import seaborn as sns

sns.set(style="white")

# Create a dataset with many short random walks
rs = np.random.RandomState(4)
pos = rs.randint(-1, 2, (10, 5)).cumsum(axis=1)
pos -= pos[:, 0, np.newaxis]
step = np.tile(range(5), 10)
walk = np.repeat(range(10), 5)
df = pd.DataFrame(np.c_[pos.flat, step, walk],
                  columns=["position", "step", "walk"])



# Initialize a grid of plots with an Axes for each walk
grid = sns.FacetGrid(df, col="walk", hue="walk", col_wrap=5, size=5,
        aspect=1)


# Draw a bar plot to show the trajectory of each random walk
grid.map(sns.barplot, "step", "position", palette="Set3").add_legend();

grid.savefig("/Users/mymacmini/Desktop/test_fig.png")
#sns.plt.show()

Which makes this plot:

As you can see I get the legend wrong. How can I make it right?

Answer 1

Some how there is one legend item for each of the subplot. Looks like if we want to have legend corresponds to the bars in each of the subplot, we have to manually make them.

# Let's just make a 1-by-2 plot
df = df.head(10)

# Initialize a grid of plots with an Axes for each walk
grid = sns.FacetGrid(df, col="walk", hue="walk", col_wrap=2, size=5,
        aspect=1)

# Draw a bar plot to show the trajectory of each random walk
bp = grid.map(sns.barplot, "step", "position", palette="Set3")

# The color cycles are going to all the same, doesn't matter which axes we use
Ax = bp.axes[0]

# Some how for a plot of 5 bars, there are 6 patches, what is the 6th one?
Boxes = [item for item in Ax.get_children()
         if isinstance(item, matplotlib.patches.Rectangle)][:-1]

# There is no labels, need to define the labels
legend_labels  = ['a', 'b', 'c', 'd', 'e']

# Create the legend patches
legend_patches = [matplotlib.patches.Patch(color=C, label=L) for
                  C, L in zip([item.get_facecolor() for item in Boxes],
                              legend_labels)]

# Plot the legend
plt.legend(handles=legend_patches)

Answer 2

When the legend doesn't work out you can always make your own easily like this:

import matplotlib

name_to_color = {
    'Expected':   'green',
    'Provided':   'red',
    'Difference': 'blue',
}

patches = [matplotlib.patches.Patch(color=v, label=k) for k,v in name_to_color.items()]
matplotlib.pyplot.legend(handles=patches)