Matplotlib绘制氢原子能级图调整图例顺序保持颜色对应问题

问题根源 #

matplotlib默认按照线条绘制的先后顺序匹配图例传入的标签序列，你直接反序图例标签但没有调整对应线条的匹配关系，就会出现颜色和标签错乱的问题。
你当前代码中绘制线条的顺序是先画n=1（红色）、再画n=2、n=3……直到n=10，原本直接传入反序的标签自然会和绘制顺序的线条错配。

解决方法 #

绘图时存储每条线的句柄和对应标签，最后将存储的序列反转再传入图例即可，不需要修改绘图逻辑和颜色匹配规则。

修改后可运行代码 #

import numpy
import math
import matplotlib.pyplot as plt

colors = ['#e6194b', '#3cb44b', '#ffe119', '#4363d8', '#f58231', '#911eb4', '#46f0f0', '#f032e6', '#bcf60c', '#fabebe', '#008080', '#e6beff', '#9a6324', '#fffac8', '#800000', '#aaffc3', '#808000', '#ffd8b1', '#000075', '#808080', '#ffffff', '#000000']
# 可自定义参数
n_int = 6 # 初始能级，必须大于n_final
n_final = 5 # 最终能级 
n_max = 11 # 最高能级+1，如要显示1-10能级则设为11

# 氢原子能级计算参数
m_e = 9.1093837015* 10 ** -31 # kg
e = 1.602176634* 10 ** -19 # C
vp = 1.113* 10 ** -10 # C^2 / Jm
h = 6.62607015 * 10 ** -34 # J*s
h_bar = h/(2 * math.pi)  
n_min = 1
n = numpy.linspace(n_min, n_max, n_max-n_min)

# 能级计算公式
p1 = m_e * e **4
p2 = 2 * vp ** 2 
p3 = h_bar ** 2
p4 = n ** 2
p5 = p2*p3*p4
e_lv = - p1/p5 # 输出输入范围内的所有能级值

# 补充原代码缺失的变量定义
e_lv_i = e_lv[n_int-1]
e_lv_f = e_lv[n_final-1]
e_lv_math = -numpy.log10(numpy.abs(e_lv_i))
e_lv_math2 = -numpy.log10(numpy.abs(e_lv_f))
d_e_math = e_lv_math2-e_lv_math

inc = 0
# 新增：存储线条句柄和对应标签
handles = []
for big_int in e_lv:
    x_big = range(0,n_max)
    y_big = [big_int]*n_max
    # 接收绘图返回的线条句柄
    line, = plt.plot(x_big, -numpy.log10(numpy.abs(y_big)), color = colors[inc])
    handles.append((line, str(inc+1)))
    inc+=1

# 以下配置不需要放在循环里重复执行，挪到外部即可
plt.xticks([])
plt.yticks([])
plt.ylabel('Relative Energy Levels')
plt.arrow(4.5, e_lv_math, 0, d_e_math, width = 0)
plt.tight_layout()

# 反转句柄序列，让n=1出现在图例最底部
handles_reversed = handles[::-1]
lines, labels = zip(*handles_reversed)
# 传入反转后的句柄和标签生成图例
plt.legend(lines, labels, loc="center right", bbox_to_anchor = (1.2,.5), title = "n=")
plt.show()

效果参考 #

原始正确对应关系的效果图 #

直接修改标签导致错乱的错误效果图 #

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