基於matplotlib的2D/3D抽象網格和能量曲線繪製程序

前言

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

• GitHub地址: PytLab/catplot
• PyPI地址: Python Package Index

正文

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

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

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

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

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

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

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

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

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

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

# 從catplot中導入繪製所需的組件: 畫布 和 線from catplot.ep_components.ep_canvas import EPCanvasfrom 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，比如拼接，合併，平移，添加陰影、改變顏色, 輔助線， 修改畫布大小，導出插值數據等等。具體的例子參考: PytLab/catplot

繪製二維和三維抽象網格

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

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

繪製5x5的二維網格

notebook版可以參見：PytLab/catplot

創建nodes和edges

from catplot.grid_components.nodes import Node2Dfrom catplot.grid_components.edges import Edge2Dnodes, edges = [], []# 創建重複單元中的nodes和edgetop = 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的edgese1 = Edge2D(top, t1, width=4)e2 = Edge2D(top, t2, width=4)edges.extend([e1, e2])# 中間的nodesbridge1 = 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])# 連接他們的edgese1 = Edge2D(bridge1, b1)e2 = Edge2D(bridge1, bridge2)e3 = Edge2D(bridge2, b2)e4 = Edge2D(b1, b2)edges.extend([e1, e2, e3, e4])# 正中間的nodeh = Node2D([0.5, 0.5], style="h", size=700, color="#5A5A5A", alpha=0.3)nodes.append(h)

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

from catplot.grid_components.grid_canvas import Grid2DCanvasfrom catplot.grid_components.supercell import SuperCell2Dcanvas = 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: PytLab/catplot

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

怎麼可能，雖然所有的坐標都是在分數坐標系中定義的，但是在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()

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

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上 (PytLab/catplot )，歡迎有需要的童鞋參考和使用。