求助：Python matplotlib绘制氢原子能级图的缩放及辨识度优化

以下是添加了局部放大框的可运行修改代码，已删除原代码中未使用的os、pandas依赖：

import numpy
import math
import matplotlib.pyplot as plt

n_max = 10 # 想要展示的最高能级
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 = numpy.linspace(1, n_max, n_max)

p1 = m_e * e **4
p2 = 2 * vp ** 2 
p3 = h_bar ** 2
p4 = n ** 2
p5 = p2*p3*p4
e_lv = - p1/p5

# 创建主画布和坐标轴
fig, ax = plt.subplots(figsize=(6,8))
# 主图绘制所有能级
for inte in e_lv:
    x = range(0,n_max)
    y = [inte]*n_max
    ax.plot(x, y)
ax.set_xticks([])
ax.set_ylabel('能量 (J)')
ax.set_title('氢原子能级图')

# ------------------- 添加局部放大框 -------------------
# 定义放大框在主图的位置：[左下角x坐标, 左下角y坐标, 宽度, 高度]，数值为相对主图的比例
ax_inset = ax.inset_axes([0.15, 0.6, 0.3, 0.35])
# 放大框内重新绘制所有能级
for inte in e_lv:
    x = range(0,n_max)
    y = [inte]*n_max
    ax_inset.plot(x, y)
# 设置放大框的y轴范围，仅展示高能级区间（这里设置为n≥5的能级区间，可根据需求调整）
ax_inset.set_ylim(e_lv[4] - 1e-20, 0)
ax_inset.set_xticks([])
# 主图上标记放大区域，自动添加连接线
ax.indicate_inset_zoom(ax_inset, edgecolor='red')

plt.show()

额外优化建议 #

• 调整inset_axes的位置参数和set_ylim的范围，即可自定义放大区域的位置和覆盖的能级区间
• 若不需要严格展示能量数值的相对比例，可直接将y轴替换为主量子数n的离散刻度，彻底解决能级堆叠问题
• 可在能级线旁添加对应n值的文本标注，提升可读性

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