實測 Python Numba JIT
4 min readAug 17, 2018
Python 的好，大家都知道；Python 的慢，大家也都知道。但就是因為如果只跑簡單的資料分析 (pandas) 或是儀器操控 (pyvisa)，計算量不大的話，Python 寫起來還是很容易上手，因此廣受歡迎、在近年來變成主流。再搭上最近的人工智慧 (深度學習) 風潮，套件如雨後春筍般冒出，Python 簡直是在江湖闖蕩的基本技能。因為它寫起來簡潔，工作效率高，所以對我這種只是偶爾寫些小程式的人來說，也是從 C++、Java，一路跳來 Python。
不過最近在跑一些多迴圈計算的時候，Python 的慢馬上就讓人很頭痛了。例如，算一個 3-dimensional tight-binding model 的 density of states，如果我們直接用數的，需要三層迴圈：
def tightbind3D(n, t, a):
result = []
Pi = 3.1415926311for i in range(n):
for j in range(n):
for k in range(n):
result.append(-2*t*(np.cos((-Pi/a)+i*(2*Pi)/(n*a)) + np.cos((-Pi/a)+j*(2*Pi)/(n*a)) + np.cos((-Pi/a)+k*(2*Pi)/(n*a))))
return result
把執行時間對 n 作圖，會得到：
對 n³ 作圖則會得到：
可以看到一個非常線性的關係；當到了 n=300 (n³=27000000) 時，跑完整個程式已經需要超過一分鐘。而為了之後的作圖，我們可能需要 n=500 甚至是 600，很明顯地運算時間會過長。
在這種情況下，numba 就能夠派上用場了。如果我們使用 numba 套件裡的 jit (just-in-time compilation) 功能，僅需對程式碼做很簡單的修改：
from numba import jit@jit(nopython=True)
def tightbind3D(n, t, a):
result = []
Pi = 3.1415926311for i in range(n):
for j in range(n):
for k in range(n):
result.append(-2*t*(np.cos((-Pi/a)+i*(2*Pi)/(n*a)) + np.cos((-Pi/a)+j*(2*Pi)/(n*a)) + np.cos((-Pi/a)+k*(2*Pi)/(n*a))))
return result
把套用了 jit 之後執行的時間和原始的程式做比較，發現速度快了 48 倍 (使用 jit 會有個 ~0.2s 的包裝時間，因此有個明顯的截距)：
另外，numba 亦支援 cuda 語言，有 gpu 的使用者也可以用 cuda Python 來加速計算。Nvidia 官方提供的範例為簡單的向量相加：
