如何根据数组值动态调整Matplotlib绘图线段颜色以对应不同情绪分类

我尝试过与我的问题类似的其他解决方案但均未成功。
我的需求是绘图线段颜色随情绪值变化：情绪值为0时显示黑色，值为1时显示红色，值为2时显示蓝色，值为3时显示黄色，值为4时显示绿色，以此类推。
我最初实现的代码如下：

import numpy as np
import matplotlib.pyplot as plt

random_emotions = [0,0,0,0,0,0,0,1,2,3,2,1,2,1,
           2,3,2,2,2,2,1,1,2,3,3,3,3,3,3,4,
           4,4,4,4,2,1,2,2,1,2,3,4,0,0,0,0,0]
random_emotions = np.array(random_emotions)

EmotionsInNumber = random_emotions

x = np.array(list(range(0,len(EmotionsInNumber))))

Angry = np.ma.masked_where(EmotionsInNumber == 0,EmotionsInNumber)
Fear = np.ma.masked_where(EmotionsInNumber == 1,EmotionsInNumber)
Happy  = np.ma.masked_where(EmotionsInNumber == 2,EmotionsInNumber)
Neutral = np.ma.masked_where(EmotionsInNumber == 3, EmotionsInNumber)
Sad = np.ma.masked_where(EmotionsInNumber == 4,EmotionsInNumber)

fig, ax = plt.subplots()
ax.plot(x, Angry,linewidth = 4, color = 'black')
ax.plot(x, Fear,linewidth = 4, color = 'red')
ax.plot(x, Happy,linewidth = 4, color = 'blue')
ax.plot(x, Neutral,linewidth = 4, color = 'yellow')
ax.plot(x, Sad,linewidth = 4, color = 'green')
ax.legend(['Angry','Fear','Happy','Neutral','Sad',])  
ax.set_title("Emotion Report of ")

plt.show()

运行后线段颜色未按预期变化，图例显示错误。另外找到的同类方案存在线段颜色重叠的问题，推测是按颜色分组的逻辑有误。

问题原因 #

1. 原掩码逻辑写反：np.ma.masked_where(条件, 数组)会把符合条件的位置设为掩码，你要保留情绪值为0的点，应该遮掉不等于0的位置，而不是等于0的位置。
2. 直接多次调用plot绘制掩码后的数组会出现线段断裂、重叠的问题，更合理的方式是用LineCollection逐段匹配颜色，实现连续的多色线段。

修正后代码 #

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# 原始数据
random_emotions = [0,0,0,0,0,0,0,1,2,3,2,1,2,1,
           2,3,2,2,2,2,1,1,2,3,3,3,3,3,3,4,
           4,4,4,4,2,1,2,2,1,2,3,4,0,0,0,0,0]
EmotionsInNumber = np.array(random_emotions)
x = np.arange(len(EmotionsInNumber))
y = EmotionsInNumber  # 这里y轴用情绪值，你也可以换成自己需要的其他数值序列

# 情绪值对应颜色映射
color_map = {
    0: 'black',   # Angry
    1: 'red',     # Fear
    2: 'blue',    # Happy
    3: 'yellow',  # Neutral
    4: 'green'    # Sad
}
label_map = {
    0: 'Angry',
    1: 'Fear',
    2: 'Happy',
    3: 'Neutral',
    4: 'Sad'
}

# 构造线段集合：每两个相邻点组成一段线段
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# 给每个线段匹配对应情绪值的颜色，用前一个点的情绪值作为线段的颜色
colors = [color_map[e] for e in EmotionsInNumber[:-1]]

# 绘制线段
fig, ax = plt.subplots()
lc = LineCollection(segments, colors=colors, linewidth=4)
ax.add_collection(lc)

# 配置坐标轴和图例
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min()-0.5, y.max()+0.5)
# 手动构造图例项
from matplotlib.lines import Line2D
legend_elements = [Line2D([0], [0], color=color_map[e], lw=4, label=label_map[e]) for e in color_map]
ax.legend(handles=legend_elements)
ax.set_title("Emotion Report of ")

plt.show()

逻辑说明 #

• 把所有相邻的两个坐标点组合成独立线段，每段线段的颜色和前一个点对应的情绪值绑定，不会出现不同颜色重叠的问题
• 手动构造图例项，避免原方案中图例和实际绘制内容不匹配的问题
• 不需要多次调用plot函数，所有线段一次性渲染，性能和显示效果都更好

内容的提问来源于stack exchange，提问作者sjsu