【matplotlib】凡例を色毎に1つのみ、あるいはカテゴリ分けして表示する方法

概要

matplotlibでグラフを作っていると、凡例を色毎に1つだけ表示したい、あるいはカテゴリに分けて表示したい場合があると思う。matplotlibではlabelという引数で凡例を管理することが多いが、ここではもう少し汎用性が高い方法を紹介する。

実装コード

Google Colabで作成した本記事のコードは、こちらにあります。

各種インポート

各種インポート

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerTuple

実行時の各種バージョン

Name	Version
numpy	1.21.6
matplotlib	3.2.2

方法

plot情報をそれぞれの変数p*に渡してlegendのオプションを指定して表示する。
handler_map={tuple: HandlerTuple(ndivide=None)}は、凡例の左側のプロットの種類の表示する分割数である。ndivide=Noneではtupleサイズに合わせるので、全てのプロットの種類が表示される。基本、ndivide=Noneで問題ないと思う。(例えばndivide=1とすると分割数が1のため、p1とp2が重なって表示される。)詳しくは、matplotlib.legend_handler.HandlerTupleを参照されたい。

複数の凡例をカテゴリに分けて表示

fig, ax = plt.subplots()
p1, = ax.plot([1, 2.5, 3], 'r-d')
p2, = ax.plot([3, 2, 1], 'k-o')
p3, = ax.plot([2, 1, 2], 'b-+')

ax.legend([(p1, p2), p3], ['Two keys', 'One key'],
           handler_map={tuple: HandlerTuple(ndivide=None)})

plt.show()

出力結果

実用例

色毎に1つのみの凡例を表示

例えば、複数の同じ性質のグラフを表示するのに色毎に1つのみの凡例を表示したい場合がある。例として、sin波がy軸方向で平行移動したグラフに1つの判例をつけてみる。

目標とする図

表示例1

listに渡す

x = np.linspace(0, 4*np.pi)
ps = [0] * 10

fig, ax = plt.subplots()
for i, shift_y in enumerate(range(10)):
    y = np.sin(x) + shift_y
    ps[i], = ax.plot(x, y, 'b')

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend([tuple(ps)], ['wave'], handler_map={tuple: HandlerTuple(ndivide=None)})
plt.show()

このように、listに1つ1つのplot情報を渡して、tupleで1つにまとめて表示できる。ただ、plot情報を全て渡すのは基本に忠実であるが手間なので、変数に再代入する次の方法も良いでしょう。

表示例2

変数に再代入

x = np.linspace(0, 4*np.pi)

fig, ax = plt.subplots()
for shift_y in range(10):
    y = np.sin(x) + shift_y
    p, = ax.plot(x, y, 'b')

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend([p], ['wave'], handler_map={tuple: HandlerTuple(ndivide=None)})
plt.show()

1つのみの凡例を表示したいケースでは、コード内でも同じ分類をしたいので、変数に再代入しても可読性が損なわないと思う。そのことを踏まえると、この方法がシンプルで良い気がします。ちなみに、matplotlibでlabelを使うなら、次のような書き方になると思います。

表示例3

labelを使う

x = np.linspace(0, 4*np.pi)

fig, ax = plt.subplots()
for i, shift_y in enumerate(range(10)):
    y = np.sin(x) + shift_y
    if i == 0:
        ax.plot(x, y, 'b', label='wave')
    else:
        ax.plot(x, y, 'b')

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend()
plt.show()

可読性は落ちますが、legend()オプション無しで凡例が表示できるのは使いやすいですね。

カテゴリ分けした凡例

例えば、複数の同じ性質のグラフを表示するのに色分けとプロットの種類でカテゴリ分けしたい場合がある。例として、元の関数とその微分した関数を、色分けとプロットの種類でカテゴリ分けして凡例に表示する。

目標とする図

カテゴリ分けした凡例

x = np.linspace(0, 2)
y1 = x ** 2
y1_prime = 2 * x
y2 = -x ** 2
y2_prime = -2 * x
y_0 = x * 0

fig, ax = plt.subplots()
y1_label, = ax.plot(x, y1, 'r')
y1_prime_label, = ax.plot(x, y1_prime, 'r:')
y2_label, = ax.plot(x, y2, 'b')
y2_prime_label, = ax.plot(x, y2_prime, 'b:')
y_0_label, = ax.plot(x, y_0, 'k--')

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend([(y1_label, y1_prime_label), (y2_label, y2_prime_label), y_0_label],
          ['y1=x**2, dy1/dx=2x', 'y2=-x**2, dy2/dx=-2x', '0'],
          handler_map={tuple: HandlerTuple(ndivide=None)})
plt.show()

まとめ

matplotlibのlabelで限界を感じたので、凡例の表示方法を公式で見つけて記事にしました。誰かのお役に立てれば幸いです。

参考資料

matplotlib Legend guide