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

Matplotlib로 산점도 추세선 그리는 방법 완벽 가이드

Matplotlib로 산점도(scatter plot)에 추세선(trend line)을 추가하려면 NumPy의 polyfit()poly1d() 메서드를 활용하면 됩니다. polyfit()은 최소제곱법(least squares)을 기반으로 데이터에 가장 잘 맞는 다항식의 계수를 계산하고, poly1d()는 이 계수를 다항식 객체로 변환하여 추세선 위의 좌표값을 손쉽게 구할 수 있게 해줍니다.

구현 단계

  • 그림(figure) 크기를 설정하고, 서브플롯 사이 및 주변 여백을 자동으로 조정합니다.
  • NumPy를 사용하여 x, y 데이터 포인트를 생성합니다.
  • 그림과 서브플롯(axes) 세트를 생성합니다.
  • 산점도를 그려 x, y 데이터 포인트를 시각화합니다.
  • polyfit()poly1d() 메서드로 추세선 데이터 포인트를 계산합니다.
  • plot() 메서드를 사용하여 x와 p(x) 값을 선 그래프로 그립니다.
  • show() 메서드로 최종 그림을 화면에 표시합니다.

예제 코드

import numpy as np
from matplotlib import pyplot as plt

# 그림 크기 및 레이아웃 설정
plt.rcParams["figure.figsize"] = [7.50, 3.50]
plt.rcParams["figure.autolayout"] = True

# 무작위 x, y 데이터 생성
x = np.random.rand(100)
y = np.random.rand(100)

fig, ax = plt.subplots()
_ = ax.scatter(x, y, c=x, cmap='plasma')

# 1차 다항식(직선)으로 추세선 계산
z = np.polyfit(x, y, 1)
p = np.poly1d(z)
plt.plot(x, p(x), "r-o")

plt.show()

코드 설명

  • np.polyfit(x, y, 1): 1차 다항식(직선) 기울기와 절편을 최소제곱법으로 구합니다. 차수를 2 이상으로 바꾸면 곡선 형태의 추세선도 그릴 수 있습니다.
  • np.poly1d(z): 계수 배열 z를 호출 가능한 다항식 함수로 변환합니다.
  • c=x, cmap='plasma': x 값에 따라 점의 색상을 다르게 표현하여 데이터 분포를 한눈에 파악하기 쉽게 합니다.

실행 결과

코드를 실행하면 plasma 컬러맵이 적용된 산점도 위에 빨간색 직선 형태의 추세선이 함께 표시됩니다. 이를 통해 두 변수 간의 전반적인 경향성을 직관적으로 확인할 수 있습니다.