Python カテゴリ別に色を指定して、散布図を作成する「JupyterLab向け」

'21/12/22更新：カテゴリの色を指定するコードを追記しました。

上図のデータ元は、下図のように横軸がA列, 縦軸がE列、カテゴリ変数はF列です。

■本プログラム

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import pandas as pd

file_path = 'test_data.csv'
df = pd.read_csv(file_path)
df


# In[2]:


DF = df.copy()
# カテゴリ変数にしたい列をカテゴリ型へ変更
DF['machine_no'] = DF['machine_no'].astype('category')
# 型をチェック
print(DF.dtypes)


# In[3]:


label = 'machine_no'
legend_list = list(DF.groupby(label).count().index)
cnt = len(legend_list)
print(legend_list, cnt)


# In[4]:


x_name = 'x1'
y_name = 'stress'

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
plt.rcParams['font.size'] = 18 # グラフの基本フォントサイズの設定

# プロットするカラーを指定するための処理
color_list = []
for j in range(cnt):
    color_list.append(cm.cool(j/cnt)) # jet cool autumn

fig = plt.figure(dpi=80, figsize=(6,4))

# ひとつの散布図にカテゴリ別にプロットする
handle_list = []
for i, legend in enumerate(DF[label].unique()):
    print(i, legend)
    df_buf = DF[DF[label] == legend]
    plt.plot(df_buf[x_name], df_buf[y_name], c=color_list[i])
    p = plt.scatter(df_buf[x_name], df_buf[y_name], c=color_list[i])
    handle_list.append(p)
    
plt.legend(
    handles = handle_list,
    labels = legend_list,
    bbox_to_anchor = (1.00, 0.9),
    title = label
)
plt.xlabel(x_name)
plt.ylabel(y_name)
plt.title('title')

x_min = 0
x_max = 20
y_spec = 150

# x軸の範囲（レンジ）を指定
plt.xlim(x_min, x_max)
# 横線を入れる
plt.hlines(y_spec, x_min, x_max, 'r', linestyles='dashed')
# テキストを挿入する
plt.text(x_min, y_spec + 10, r'Spec 150[MPa]', size=18, color='r')
plt.grid()
plt.tick_params()
#plt.tight_layout()
#plt.show()
plt.savefig(file_path[:-4] + '.jpg', bbox_inches="tight")

以上

<広告>

リンク

リンク