用matplotlib做能级图

Posted on | Edited on

最近因为交作业的需要，要做一张研究分子或者原子的电子结构时非常常见的能级图（如下图，来自网络）。

这张图看着简单，实际做起来还真没有找到现成的工具。现在我手里有一堆能级float数据存在一维numpy数组e_levels里，怎样才能做出类似这样的图呢？

直方图

能级图本身研究的是一维数据的分布，因此我首先想到的是matplotlib的直方图，代码也简单:

1
2
plt.hist(e_levels, bins=3000)
plt.show()

因为需要尽可能看到细粒度的能级分布，所以需要用非常多的bins，作图时间比较长，效果如下：

可以看到大概的意思有了，但是非常丑。一方面能级不是垂直分布的，另一方面相近的能级还是会分到一个bin里去，导致出现有纵向长度不一样的情况。

rugplot

seaborn是matplotlib的一个高级包装，有许多实用的绘图函数。我在翻查文档时发现了一个“貌似”很对口的函数rugplot，中文译名不知为何。rugplot和直方图不同，摆明了就是要画细线，还可以自定义细线的高度（画多长），绘图轴等。实际生活中，能找到这样的满足需求的现成代码，是解决问题效率最高的方式。

1
2
3
fig = plt.figure(figsize=(2,6)) # 使能级图具有较高的高度和较小的宽度
sns.rugplot(e_levels, height=1, ax=plt.gca(), axis='y') # height为所绘制细线相对图大小的高度
plt.show()


有那么一点意思，但也有一个很关键的问题：虽然我们通过定义height让每根线都画的很长（height默认只有0.05)，但每根线都直接横跨整个x轴，有点业余，将来如果在y轴上加标注，会显得更加难看。
下面这两张图，分别对应的是把height设置成0.8的情况和添加一行ax.set_xlim(-0.2, 0.2)使x轴变宽的情况，都没有达到我们期望的，有几根较长的线居中哪怕靠右排布的目的。

自己动手，丰衣足食

试了别人的轮子，github上也找了一圈无果，现在该自己动手了。前面一直没有提到，其实我还有“在图中突出某一个或多个能级”的需求，这几乎只有通过自己动手的可定制性才能达到。
首先研究了一下，matplotlib上是怎么画线的。第一步是定义一个Line2D对象：

1
2
from matplotlib import lines
line = lines.Line2D((0, 0.5), (0, 0.5)) # 参数为((x1, x2), (y1, y2))定义平面上的一个线段

然后在一个axes上添加这个对象，全部代码如下：

1
2
3
4
5
6
from matplotlib import lines
line = lines.Line2D((0, 0.5), (0, 0.5))
fig = plt.figure()
ax = plt.gca()
ax.add_line(line)
plt.show()


根据这种添加线段的办法，可以把e_levels中的float数据变成水平线段：

1
2
3
4
5
6
7
8
9
10
padding = e_levels.ptp() / 10 # 首先为能级图确定一个边距，以免最高能级和最低能级和边框相距过近
figure = plt.figure(figsize=(2,6))
ax = plt.gca()
ax.get_xaxis().set_visible(False) # x轴无实际意义，隐藏
ax.get_yaxis().set_label_text('Energy')
ax.set_ylim(e_levels.min() - padding, e_levels.max() + padding) # 根据边距设定y轴范围
for eval_ in e_levels:
    line = lines.Line2D((0.3, 0.7), (eval_, eval_), c='black') # 逐个添加能级
    ax.add_line(line)
plt.show()


可以看到基本达到了目的。由于这张图是自己一点点画出来的，可定制性很强，可以针对某个能级调整不同的颜色、宽度、透明度，或者加文字标注等等。比如，根据matplotlib官方文档，可以进行如下设置启用\(\LaTeX\)文字渲染：

1
2
3
4
5
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)

然后在我们画能级的代码中加入这样一段，就可以为我所关心的第0、49、50、51和100个能级做一个简单的标注：

1
2
3
4
for idx in [0, 49, 50, 51, 100]: # 所关心能级的index
    # 函数原型ax.text(x, y, s, fontdict=None, withdash=False, **kwargs)
    # x=0.72由划线横坐标到0.7微调而来
    ax.text(x=0.72, y=e_levels[idx]-0.07, s=r'$\varepsilon_{' + str(idx) + '}$', fontsize=14)

最后绘图的结果就像这样；

后记

辨析一组单词：axis和axes

axes是axis的复数形式。在matplotlib（以及matlab）中，axis顾名思义就是一个坐标轴，而axes从字面意义上看是多个坐标轴，实际操作的意义是一张作的二维或者三维图。当然，这里的“图”也不是figure，通过subplot，一个figure可以包含多个axes。
这其实给以中文为母语的Coder带来了混淆，因为我们无法给axes做一个好的翻译。从axis引申到axes成为一个“图表”，从英语的角度来看是容易理解的，但汉语中从“坐标轴”引申到“多坐标轴”就让人摸不着头脑。其结果是不论是axes还是figure我们都通通叫做“图”，figure->axes->axis的清晰包含关系不复存在，为新人入门平添了困难。

图表的label在savefig时遭到截断

我在第一次试图保存能级图时，y轴的标注“Energy”惨遭截断（尽管plt.show()看到的完全正常）

参考stackoverflow得到了解决。简单处理的话，只需加入以下代码即可：

1
2
from matplotlib import rcParams
rcParams.update({'figure.autolayout': True})