簡介
通常來說,Python不是一種高效能的語言,在某種意義上,這種說法是真的。但是,隨著以Numpy為中心的數學和科學軟體包的生態圈的發展,達到合理的效能不會太困難。
當效能成為問題時,執行時間通常由幾個函式決定。用C重寫這些函式,通常能極大的提升效能。
在本系列的第一部分中,我們來看看如何使用NumPy的C API來編寫C語言的Python擴充套件,以改善模型的效能。在以後的文章中,我們將在這裡提出我們的解決方案,以進一步提升其效能。
檔案
這篇文章中所涉及的檔案可以在Github上獲得。
模擬
作為這個練習的起點,我們將在像重力的力的作用下為N體來考慮二維N體的模擬。
以下是將用於儲存我們世界的狀態,以及一些臨時變數的類。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# lib/sim.py class World(object): """World is a structure that holds the state of N bodies and additional variables. threads : (int) The number of threads to use for multithreaded implementations. STATE OF THE WORLD: N : (int) The number of bodies in the simulation. m : (1D ndarray) The mass of each body. r : (2D ndarray) The position of each body. v : (2D ndarray) The velocity of each body. F : (2D ndarray) The force on each body. TEMPORARY VARIABLES: Ft : (3D ndarray) A 2D force array for each thread's local storage. s : (2D ndarray) The vectors from one body to all others. s3 : (1D ndarray) The norm of each s vector. NOTE: Ft is used by parallel algorithms for thread-local storage. s and s3 are only used by the Python implementation. """ def __init__(self, N, threads=1, m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e-3): self.threads = threads self.N = N self.m = np.random.uniform(m_min, m_max, N) self.r = np.random.uniform(-r_max, r_max, (N, 2)) self.v = np.random.uniform(-v_max, v_max, (N, 2)) self.F = np.zeros_like(self.r) self.Ft = np.zeros((threads, N, 2)) self.s = np.zeros_like(self.r) self.s3 = np.zeros_like(self.m) self.dt = dt |
在開始模擬時,N體被隨機分配質量m,位置r和速度v。對於每個時間步長,接下來的計算有:
- 合力F,每個體上的合力根據所有其他體的計算。
- 速度v,由於力的作用每個體的速度被改變。
- 位置R,由於速度每個體的位置被改變。
第一步是計算合力F,這將是我們的瓶頸。由於世界上存在的其他物體,單一物體上的力是所有作用力的總和。這導致複雜度為O(N^2)。速度v和位置r更新的複雜度都是O(N)。
如果你有興趣,這篇維基百科的文章介紹了一些可以加快力的計算的近似方法。
純Python
在純Python中,使用NumPy陣列是時間演變函式的一種實現方式,它為優化提供了一個起點,並涉及測試其他實現方式。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
# lib/sim.py def compute_F(w): """Compute the force on each body in the world, w.""" for i in xrange(w.N): w.s[:] = w.r - w.r[i] w.s3[:] = (w.s[:,0]**2 + w.s[:,1]**2)**1.5 w.s3[i] = 1.0 # This makes the self-force zero. w.F[i] = (w.m[i] * w.m[:,None] * w.s / w.s3[:,None]).sum(0) def evolve(w, steps): """Evolve the world, w, through the given number of steps.""" for _ in xrange(steps): compute_F(w) w.v += w.F * w.dt / w.m[:,None] w.r += w.v * w.dt |
合力計算的複雜度為O(N^2)的現象被NumPy的陣列符號所掩蓋。每個陣列操作遍歷陣列元素。
視覺化
這裡是7個物體從隨機初始狀態開始演化的路徑圖:
效能
為了實現這個基準,我們在專案目錄下建立了一個指令碼,包含如下內容:
1 2 3 |
import lib w = lib.World(101) lib.evolve(w, 4096) |
我們使用cProfile模組來測試衡量這個指令碼。
1 |
python -m cProfile -scum bench.py |
前幾行告訴我們,compute_F確實是我們的瓶頸,它佔了超過99%的執行時間。
1 2 3 4 5 6 7 8 9 10 11 |
428710 function calls (428521 primitive calls) in 16.836 seconds Ordered by: cumulative time ncalls tottime percall cumtime percall filename:lineno(function) 1 0.000 0.000 16.837 16.837 bench.py:2(<module>) 1 0.062 0.062 16.756 16.756 sim.py:60(evolve) 4096 15.551 0.004 16.693 0.004 sim.py:51(compute_F) 413696 1.142 0.000 1.142 0.000 {method 'sum' ... 3 0.002 0.001 0.115 0.038 __init__.py:1(<module>) ... |
在Intel i5桌上型電腦上有101體,這種實現能夠通過每秒257個時間步長演化世界。
簡單的C擴充套件 1
在本節中,我們將看到一個C擴充套件模組實現演化的功能。當看完這一節時,這可能幫助我們獲得一個C檔案的副本。檔案src/simple1.c
,可以在GitHub上獲得。
關於NumPy的C API的其他文件,請參閱NumPy的參考。Python的C API的詳細文件在這裡。
樣板
檔案中的第一件事情是先宣告演化函式。這將直接用於下面的方法列表。
1 |
static PyObject *evolve(PyObject *self, PyObject *args); |
接下來是方法列表。
1 2 3 4 |
static PyMethodDef methods[] = { { "evolve", evolve, METH_VARARGS, "Doc string."}, { NULL, NULL, 0, NULL } /* Sentinel */ }; |
這是為擴充套件模組的一個匯出方法列表。這隻有一個名為evolve方法。
樣板的最後一部分是模組的初始化。
1 2 3 4 |
PyMODINIT_FUNC initsimple1(void) { (void) Py_InitModule("simple1", methods); import_array(); } |
另外,正如這裡顯示,initsimple1中的名稱必須與Py_InitModule
中的第一個引數匹配。對每個使用NumPy API的擴充套件而言,呼叫import_array
是有必要的。
陣列訪問巨集
陣列訪問的巨集可以在陣列中被用來正確地索引,無論陣列被如何重塑或分片。這些巨集也使用如下的程式碼使它們有更高的可讀性。
1 2 3 4 5 6 7 8 |
#define m(x0) (*(npy_float64*)((PyArray_DATA(py_m) + \ (x0) * PyArray_STRIDES(py_m)[0]))) #define m_shape(i) (py_m->dimensions[(i)]) #define r(x0, x1) (*(npy_float64*)((PyArray_DATA(py_r) + \ (x0) * PyArray_STRIDES(py_r)[0] + \ (x1) * PyArray_STRIDES(py_r)[1]))) #define r_shape(i) (py_r->dimensions[(i)]) |
在這裡,我們看到訪問巨集的一維和二維陣列。具有更高維度的陣列可以以類似的方式被訪問。
在這些巨集的幫助下,我們可以使用下面的程式碼迴圈r:
1 2 3 4 5 |
for(i = 0; i < r_shape(0); ++i) { for(j = 0; j < r_shape(1); ++j) { r(i, j) = 0; // Zero all elements. } } |
命名標記
上面定義的巨集,只在匹配NumPy的陣列物件定義了正確的名稱時才有效。在上面的程式碼中,陣列被命名為py_m
和py_r
。為了在不同的方法中使用相同的巨集,NumPy陣列的名稱需要保持一致。
計算力
特別是與上面五行的Python程式碼相比,計算力陣列的方法顯得頗為繁瑣。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
static inline void compute_F(npy_int64 N, PyArrayObject *py_m, PyArrayObject *py_r, PyArrayObject *py_F) { npy_int64 i, j; npy_float64 sx, sy, Fx, Fy, s3, tmp; // Set all forces to zero. for(i = 0; i < N; ++i) { F(i, 0) = F(i, 1) = 0; } // Compute forces between pairs of bodies. for(i = 0; i < N; ++i) { for(j = i + 1; j < N; ++j) { sx = r(j, 0) - r(i, 0); sy = r(j, 1) - r(i, 1); s3 = sqrt(sx*sx + sy*sy); s3 *= s3 * s3; tmp = m(i) * m(j) / s3; Fx = tmp * sx; Fy = tmp * sy; F(i, 0) += Fx; F(i, 1) += Fy; F(j, 0) -= Fx; F(j, 1) -= Fy; } } } |
請注意,我們使用牛頓第三定律(成對出現的力大小相等且方向相反)來降低內環範圍。不幸的是,它的複雜度仍然為O(N^2)。
演化函式
該檔案中的最後一個函式是匯出的演化方法。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
static PyObject *evolve(PyObject *self, PyObject *args) { // Declare variables. npy_int64 N, threads, steps, step, i; npy_float64 dt; PyArrayObject *py_m, *py_r, *py_v, *py_F; // Parse arguments. if (!PyArg_ParseTuple(args, "ldllO!O!O!O!", &threads, &dt, &steps, &N, &PyArray_Type, &py_m, &PyArray_Type, &py_r, &PyArray_Type, &py_v, &PyArray_Type, &py_F)) { return NULL; } // Evolve the world. for(step = 0; step< steps; ++step) { compute_F(N, py_m, py_r, py_F); for(i = 0; i < N; ++i) { v(i, 0) += F(i, 0) * dt / m(i); v(i, 1) += F(i, 1) * dt / m(i); r(i, 0) += v(i, 0) * dt; r(i, 1) += v(i, 1) * dt; } } Py_RETURN_NONE; } |
在這裡,我們看到了Python引數如何被解析。在該函式底部的時間步長迴圈中,我們看到的速度和位置向量的x和y分量的顯式計算。
效能
C版本的演化方法比Python版本更快,這應該不足為奇。在上面提到的相同的i5桌上型電腦中,C實現的演化方法能夠實現每秒17972個時間步長。相比Python實現,這方面有70倍的提升。
觀察
注意,C程式碼一直保持儘可能的簡單。輸入引數和輸出矩陣可以進行型別檢查,並分配一個Python裝飾器函式。刪除分配,不僅能加快處理,而且消除了由Python物件不正確的引用計數造成的記憶體洩露(或更糟)。
下一部分
在本系列文章的下一部分,我們將通過發揮C-相鄰NumPy矩陣的優勢來提升這種實現的效能。之後,我們來看看使用英特爾的SIMD指令和OpenMP來進一步推進。
如果您有任何疑問,意見,建議或更正,請通過聯絡連結告訴我。