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

时间:2015-08-28 00:42:46

标签: python pandas matplotlib seaborn

我有以下代码:

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()

这使得这个情节:

enter image description here

正如你所看到的,我认为这个传奇是错误的。我怎样才能做对吗?

2 个答案:

答案 0 :(得分:11)

每个子图中有一个传奇项目。看起来如果我们想让图例对应每个子图中的条形图,我们必须手动制作它们。

# 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)

enter image description here

答案 1 :(得分:2)

当图例无法解决时,您总是可以这样轻松地制作自己的图片:

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)