基於 matplotlib 的抽象網格和能量曲線繪製程式

前言

最近在寫文章需要繪製一些一維的能量曲線(energy profile)和抽象的二維和三維的網格來表示晶體用來描述自己的演算法，於是自己在之前的指令碼的基礎上進行了整改寫成了只提供介面的Python庫,基本思想就是封裝了matplotlib中相關介面，方便快速搭建和定製自己的能量曲線和網格結構, 程式碼託管在GitHub上並上傳至PyPI。對於研究晶體材料的同學如果想通過python來繪製簡單的晶格影象可以參考一下。

GitHub地址:https://github.com/PytLab/catplot/

PyPI地址:https://pypi.python.org/pypi/catplot/

正文

首先還是介紹一下這個程式的用途，目前主要是提供三個主要的模組來繪製三方面的內容：

1. 繪製抽象的二維網格結構

catplot提供了豐富的介面用來定製所需要的任何二維網格並進行週期性擴充套件，如下圖是一個通過當個重複單元擴充套件出來的抽象(100)晶面的二維網格結構:

​

2. 繪製抽象的三維網格結構

同理只不過這次是在三維畫布中進行繪製並進行重複單元的週期性擴充套件，擴充套件的效果如下圖:

​

3. 通過插值演算法實現繪製”順滑”的energy profile

​

實現過程基本是通過對matplotlib提供的繪圖元件和介面進一步封裝成可以快速搭建上面三個型別影象的元件。

採用二次插值結合樣條插值方法繪製 energy profile

energy profile可以理解成在勢能面(Potential Energy Surface)上沿著某個特定的方向(反應座標方向)上能量的變化，

下面我就上一個簡單的例子來畫一條順滑的energy profile, 更多具體的例子我都已經jupyter notebook的形式放在的github上(https://github.com/PytLab/catplot/tree/master/examples)

# 從catplot中匯入繪製所需的元件: 畫布 和 線

from catplot.ep_components.ep_canvas import EPCanvas

from catplot.ep_components.ep_lines import ElementaryLine

# 建立一個用於繪製energy profile的畫布

canvas = EPCanvas()

# 建立一條能量曲線，提供的三個值分別是三個狀態下的能量數值

line = ElementaryLine([0.0, 1.2, 0.8])

# 將這條線新增到畫布中

canvas.add_line(line)

# 繪製

canvas.draw()

canvas.figure.show()

插值方法

為了能將能量最高點沿著橫座標任意位置移動，我先將頂點的兩邊用二次函式進行插值，獲取兩個不同的二次函式形式，然後根據二次函式的形式在左右兩邊插上5個點，為了能讓分開插值的兩部分看起來連續，在將上面的10個新插的點和之前的3個點進行一次spline插值即可。

# 頂點兩側進行二次插值的演算法

def quadratic_connect_interp(x1, y1, x2, y2):

    A = np.matrix([[x1**2, x1, 1],

                   [x2**2, x2, 1],

                   [2*x2, 1, 0]])

    b = np.matrix([[y1], [y2], [0]])

    x = A.I * b

    x.shape = (1, -1)

    a, b, c = x.tolist()[0]

    poly_func = lambda x: a*x**2 + b*x + c

    return poly_func

​

與插值相關的方法參考:https://github.com/PytLab/catplot/blob/master/catplot/interpolate.py

豐富的介面

除了上面最簡單的例子，catplot還提供了豐富的介面來定製和操作energy profile，比如拼接，合併，平移，新增陰影、改變顏色, 輔助線， 修改畫布大小，匯出插值資料等等。具體的例子參考:https://github.com/PytLab/catplot/tree/master/examples

​

繪製二維和三維抽象網格

晶格中的原子和鍵在catplot中被抽象成圖中的node和edge，這樣我們就可以通過建立圖中的node和edge的方式搭建我們網格的重複單元，之後可以通過重複單元的擴充套件方法來將其擴充套件成nxn或者nxnxn的網格。

實現的基本方法就是通過matplotlib提供的Line2D, Arrow和scatter相關的介面來將相應node和edge的資料新增到maptlotlib的二維或者三維畫布中然後進行繪製和顯示。下面給分別給出兩個繪製正交網格的繪製方法:

繪製5×5的二維網格

notebook版可以參見：https://github.com/PytLab/catplot/blob/master/examples/grid_2d_examples/expand_supercell.ipynb

建立nodes和edges

from catplot.grid_components.nodes import Node2D

from catplot.grid_components.edges import Edge2D

nodes, edges = [], []

# 建立重複單元中的nodes和edge

top = Node2D([0.0, 0.0], size=800, color="#2A6A9C")

t1 = Node2D([0.0, 1.0])

t2 = Node2D([1.0, 0.0])

nodes.append(top)

# 連結這三個node的edges

e1 = Edge2D(top, t1, width=4)

e2 = Edge2D(top, t2, width=4)

edges.extend([e1, e2])

# 中間的nodes

bridge1 = Node2D([0.0, 0.5], style="s", size=600, color="#5A5A5A", alpha=0.6)

bridge2 = Node2D([0.5, 0.0], style="s", size=600, color="#5A5A5A", alpha=0.6)

b1 = bridge1.clone([0.5, 0.5])

b2 = bridge2.clone([0.5, 0.5])

nodes.extend([bridge1, bridge2])

# 連線他們的edges

e1 = Edge2D(bridge1, b1)

e2 = Edge2D(bridge1, bridge2)

e3 = Edge2D(bridge2, b2)

e4 = Edge2D(b1, b2)

edges.extend([e1, e2, e3, e4])

# 正中間的node

h = Node2D([0.5, 0.5], style="h", size=700, color="#5A5A5A", alpha=0.3)

nodes.append(h)

好了，現在我們就建立一個重複單元中的所需的所有元素，可以繪製一下看看效果了

from catplot.grid_components.grid_canvas import Grid2DCanvas

from catplot.grid_components.supercell import SuperCell2D

canvas = Grid2DCanvas()

# 將上面的元素放到supercell中，後面我們將一supercell為單位進行展開

supercell = SuperCell2D(nodes, edges)

# 繪製效果

canvas.add_supercell(supercell)

canvas.draw()

canvas.figure

​

OK, 重複單元已經搭建成功，可以以他為單位進行擴充套件了, 下面我們將其沿著x和y軸方向各進行5次重複擴充套件。

# 很簡單，就一行程式碼

expanded_supercell = supercell.expand(5, 5)

來看看效果:

canvas_big = Grid2DCanvas(figsize=(30, 20), dpi=60)  # 定製畫布大小

canvas_big.add_supercell(expanded_supercell)

canvas_big.draw()

canvas_big.figure

​

是不是很直觀和簡單呢？

繪製三維網格

繪製三維網格，catplot中我都寫了與二維繪製中相對應的類和介面，這裡就不贅述了，可以參考專案中的examples:https://github.com/PytLab/catplot/tree/master/examples/grid_3d_examples/expand_3d_supercell.ipynb

是不是隻能畫正交的網格？

怎麼可能，雖然所有的座標都是在分數座標系中定義的，但是在SuperCell類中我新增了分數座標到笛卡爾座標的轉化，從而可以使得catplot繪製任意的網格。來個例子就知道了:

# 建立nodes和edges的程式碼與上面的部分完全相同

...

# 但是我們在定義supercell的時候可以修改cell_vectors引數來是重複單元發生形變

supercell = SuperCell2D(nodes, edges, cell_vectors=[[1.0, 0.0],

                                                    [0.5, 1.0]])

canvas.add_supercell(supercell)

canvas.draw()

canvas.figure.show()

來我們看看這時候的重複單元是什麼樣子:

​

然後我們再將其進行一次3×3的擴充套件看看

expanded_supercell = supercell.expand(3, 3)

canvas_big = Grid2DCanvas(figsize=(30, 20), dpi=60)

canvas_big.add_supercell(expanded_supercell)

canvas_big.draw()

canvas_big.figure.show()

​

所以基本上現在所有型別的晶格都可以通過CatPlot來繪製了。

總結

本來catplot這個庫最初是自己用matplotlib來繪圖的小指令碼，由於現在寫論文的情況下需要靈活的繪製網格圖，所以進行了重寫，現在寫成了一個封裝了matplotlib的python庫方便使用者可以快速搭建自己想要的網格圖和繪製漂亮的energy profile。程式碼和具體使用的notebook格式的例子均開源並放到了github上 (https://github.com/PytLab/catplot)，歡迎有需要的童鞋參考和使用。