Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Matplotlib에서 범례(Legend) 자동 생성하는 방법


Matplotlib 범례 자동 생성 개요

Matplotlib에서 산점도(scatter plot)를 그릴 때, 데이터의 색상이나 크기가 여러 값으로 구성되어 있으면 각각의 의미를 보여주는 범례가 필요합니다. PathCollection 객체가 제공하는 legend_elements() 메서드를 활용하면 이러한 범례를 손쉽게 자동 생성할 수 있습니다.

Matplotlib에서 범례 생성을 자동화하려면 다음 단계를 따릅니다.

  • 그림(figure) 크기를 설정하고 서브플롯 사이 및 주변의 여백을 조정합니다.
  • 샘플 데이터 개수를 나타내는 변수 N을 초기화합니다.
  • numpy를 사용하여 x, y, c, s 데이터를 생성합니다.
  • subplots() 메서드로 figure와 subplot 집합을 만듭니다.
  • 서로 다른 색상(c)과 크기(s)로 x, y 데이터 포인트를 산점도로 시각화합니다.
  • legend() 메서드로 축(ax)에 첫 번째 범례를 배치합니다.
  • add_artist() 메서드로 해당 범례를 figure에 추가합니다.
  • PathCollection을 위한 범례 핸들(handle)과 레이블(label)을 생성합니다.
  • 마커 크기에 대한 두 번째 범례를 축에 배치합니다.
  • show() 메서드로 최종 그림을 화면에 표시합니다.

예제 코드

import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["figure.figsize"] = [7.50, 3.50]
plt.rcParams["figure.autolayout"] = True

N = 45
x, y = np.random.rand(2, N)
c = np.random.randint(1, 5, size=N)
s = np.random.randint(10, 220, size=N)

fig, ax = plt.subplots()
scatter = ax.scatter(x, y, c=c, s=s)

# 색상(클래스) 기준 범례 생성
legend1 = ax.legend(*scatter.legend_elements(), loc="lower left", title="Classes")
ax.add_artist(legend1)

# 마커 크기 기준 범례 생성
handles, labels = scatter.legend_elements(prop="sizes", alpha=0.6)
legend2 = ax.legend(handles, labels, loc="upper right", title="Sizes")

plt.show()

주요 코드 설명

legend_elements()

legend_elements()는 산점도 객체(PathCollection)로부터 범례에 필요한 핸들과 레이블을 자동으로 추출하는 핵심 메서드입니다. 인자 없이 호출하면 색상(c) 값을 기준으로 클래스별 범례 요소를 만들고, prop="sizes" 옵션을 지정하면 마커 크기(s)를 기준으로 범례 요소를 생성합니다.

add_artist()

하나의 축에는 기본적으로 하나의 범례만 유지되기 때문에, 두 번째 범례를 추가하기 전에 첫 번째 범례를 ax.add_artist(legend1)로 figure에 고정해야 두 범례가 모두 화면에 함께 표시됩니다.

출력 결과

위 코드를 실행하면 왼쪽 하단에는 "Classes"(색상 기반 클래스) 범례가, 오른쪽 상단에는 "Sizes"(마커 크기) 범례가 동시에 표시된 산점도를 확인할 수 있습니다.

Matplotlib에서 범례(Legend) 자동 생성하는 방법