首页 文章

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

提问于
浏览
0

下面的代码实现了我想要做的事情,但是以非常迂回的方式实现 . 我已经四处寻找一种简洁的方法来为一个包含多个子图的图形生成单个图例,这些子图考虑了他们的标签,但无济于事 . 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')

Example plot with 9 lines and 3 categories. Each category has a different colour

1 回答

  • 2

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

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

    如果类别尚不存在,则只向字典添加句柄 .

相关问题