matplotlibで条件によって色分けをして凡例を追加したい

Question

前提・実現したいこと

　機械学習手法のロジスティック回帰と機械学習ライブラリsklearnを使って、精度判定や影響度を出力しています。
matplotlibを使って、影響度が負の時を赤色でグラフ表示、正の時青色でグラフ表示するところまでできました。
更にやりたいことは、図3のグラフから赤の時故障、青の時正常、と凡例を追加することです。
解決方法があればご回答お願いします。

発生している問題

グラフが赤の時故障、青の時正常と凡例に追加したいが、うまくいっていない

該当のソースコード

print("*** LogisticRegression.pyの実行 ***")
print("Step1. ライブラリのインポート")

# ***** ライブラリのインポート *****
import warnings
# 余分なワーニングを非表示にする
warnings.filterwarnings('ignore')

import pickle
import os
import japanize_matplotlib
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
import csv
import pandas as pd


# ***** 入力ファイルの読み込み *****
print("Step2. 入力ファイルの読み込み")
INPUT_FILE = os.path.join(os.getcwd(), "input_file")
OUTPUT_FILE = os.path.join(os.getcwd(), "output_file")
os.makedirs(OUTPUT_FILE, exist_ok=True)

THRESH = 0.3  # 故障確率をいくつ以上を故障と判定するか指定
df = pd.read_csv(os.path.join(INPUT_FILE, 'train.csv'))
df.head() #図1

print("Step3. 入力データ前処理")
# 製品故障
df['状態(予測対象)'] = df['状態(予測対象)'].map({'故障': 0, '正常': 1})
# 購入の経過月数 131.0month -> 131.0
df['購入からの経過月数'] = df['購入からの経過月数'].str.replace('month', '').astype(float)

# 機器タイプ
m_type = pd.get_dummies(df['機器タイプ'], drop_first=True, prefix='機器')
# 保守担当チーム
team=pd.get_dummies(df['保守担当チーム'],drop_first=True,prefix='チーム')
# 表の統合
df_tmp=df.drop(['機器タイプ','保守担当チーム'],axis=1)
df_merge=pd.concat([df_tmp,m_type,team],axis=1)
df_merge.head() #図2

col = ['機器_B','機器_C','チーム_Team1-2','チーム_Team2-1','チーム_Team2-2','購入からの経過月数', '稼働時平均温度',
       '稼働時平均湿度', '油圧メーター値']
x = df_merge[col]
t = df_merge['状態(予測対象)']
#訓練データ標準化
sc=StandardScaler()
new=sc.fit_transform(x)

# ***** 学習 *****
print("Step4. モデル学習")
x_train, x_test, y_train, y_test = train_test_split(new, t,
                                                    test_size=0.2, random_state=0)

#学習
model = LogisticRegression(random_state=0, C=0.1,
                           multi_class="auto", solver="lbfgs")
model.fit(x_train, y_train)
print(f'訓練データ件数{len(y_train)} 検証データ件数{len(y_test)}')

# ***** 精度評価 *****
print("Step5. 精度評価")
#正解率
score = model.score(x_train, y_train)
score2 = model.score(x_test, y_test)
print(f"Train {score:.2%}")
print(f"Test {score2:.2%}")

#影響度を知る
feature=pd.DataFrame(model.coef_[0],col)
feature.to_csv(os.path.join(OUTPUT_FILE,"feature.csv"),header=False,index=True,encoding="shift-jis")

#影響度をグラフ化する
plt.figure(figsize=(15,5))
plt.title("故障の影響度")
plt.xlabel("要因")
plt.ylabel('故障/正常の影響度')
#影響度が正の時青,影響度が負の時赤にする
color=[('r' if model.coef_[0][i]<0 else 'b') for i in range(len(model.coef_[0]))]
#影響度を絶対値にして表示
plt.bar(col,np.abs(model.coef_[0]),width=0.5,color=color)　図3


'''
影響度が正の時正常、負の時故障と凡例に表示させたい


'''

#plt.legend(loc='upper center')
plt.show()

図1 入力ファイル

図2 カテゴリ変数をダミー変数化した

図3 影響度をグラフ出力したもの
ここから凡例を追加したい

試したこと

以下のコードで場合分けしてみたが、凡例が期待通りの動作にならなかった。
図4では、凡例が9個も出てしまっているが、凡例を赤と青の2つだけにしたい。

plt.figure(figsize=(15,5))
for i in range(len(model.coef_[0])):
    if model.coef_[0][i]<0:
        plt.bar(col[i],model.coef_[0][i],width=0.5,color="r",label="故障")
    else:
        plt.bar(col[i],model.coef_[0][i],width=0.5,color="b",label="正常")
plt.legend(loc='upper left')
plt.show() #図4

図4 凡例失敗図

補足情報（FW/ツールのバージョンなど）

使用OS windows10
pythonヴァージョン python 3.9.7 32bit環境
使用ライブラリ
sklearn, matplotlib, pandas

参考
ロジスティック回帰分析で特徴量の重要度を知る方法
(https://teratail.com/questions/263204)

Answer 1

データがないので適当な数値を入れています。

import matplotlib.pyplot as plt
import matplotlib
import numpy as np
matplotlib.rc('font', family='BIZ UDGothic')

col = ['機器_B','機器_C','チーム_Team1-2','チーム_Team2-1','チーム_Team2-2','購入からの経過月数', '稼働時平均温度',
       '稼働時平均湿度', '油圧メーター値']

coef = np.array([-0.06,-0.06,0.2,0.2,0.28,-0.8,-0.31,-0.1,0.7])

plt.figure(figsize=(15,5))
coef_p = np.where(coef>=0, coef, 0)
coef_n = np.where(coef<0, -coef, 0)

plt.bar(col,coef_p,width=0.5, label='positive')
plt.bar(col,coef_n,width=0.5, label='negative',bottom=coef_p)
plt.legend()
plt.show()

実行結果

Comments

[2017yamaguchi]

わかりやすい回答ありがとうございました。
np.whereの使い方を今まで知らなかったので勉強になりました。
無事解決することができました。