Matplotlib：使用子图线标签的所有子图的自动彩色图例

Question

下面的代码实现了我想要做的事情，但是以非常迂回的方式实现了。我已经四处寻找一种简洁的方法来为包含多个子图的图形生成一个图例，该图包含了他们的标签，但无济于事。 plt.figlegend()要求您传入标签和行，而plt.legend()只需要句柄（稍好一些）。

下面我的例子说明了我想要的。我有9个向量，每个向量有3个类别之一。我想在一个单独的子图上绘制每个矢量，标记它，并绘制一个图例，表示（使用颜色）标签的含义;这是单个图上的自动行为。

您知道更好的方法来实现下面的情节吗？

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

nr_lines = 9
nr_cats = 3
np.random.seed(1337)

# Data
X = np.random.randn(nr_lines, 100)
labels = ['Category {}'.format(ii) for ii in range(nr_cats)]
y = np.random.choice(labels, nr_lines)

# Ideally wouldn't have to manually pick colours
clrs = matplotlib.rcParams['axes.prop_cycle'].by_key()['color']
clrs = [clrs[ii] for ii in range(nr_cats)]
lab_clr = {k: v for k, v in zip(labels, clrs)}

fig, ax = plt.subplots(3, 3)
ax = ax.flatten()
for ii in range(nr_lines):
    ax[ii].plot(X[ii,:], label=y[ii], color=lab_clr[y[ii]])

lines = [a.lines[0] for a in ax]
l_labels = [l.get_label() for l in lines]

# the hack - get a single occurance of each label
idx_list = [l_labels.index(lab) for lab in labels]
lines_ = [lines[idx] for idx in idx_list]
#l_labels_ = [l_labels[idx] for idx in idx_list]
plt.legend(handles=lines_, bbox_to_anchor=[2, 2.5])
plt.tight_layout()
plt.savefig('/home/james/Downloads/stack_figlegend_example.png',
            bbox_inches='tight')

Answer 1

您可以使用字典来使用标签作为键来收集它们。例如：

handles = {}

for ii in range(nr_lines):
    l1, = ax[ii].plot(X[ii,:], label=y[ii], color=lab_clr[y[ii]])

    if y[ii] not in handles:
        handles[y[ii]] = l1

plt.legend(handles=handles.values(), bbox_to_anchor=[2, 2.5])

如果类别尚未存在，您只需在字典中添加句柄。