# 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 |
在開始模擬時(shí)，N體被隨機(jī)分配質(zhì)量m，位置r和速度v。對(duì)于每個(gè)時(shí)間步長(zhǎng)，接下來的計(jì)算有：
合力F，每個(gè)體上的合力根據(jù)所有其他體的計(jì)算。
速度v，由于力的作用每個(gè)體的速度被改變。
位置R，由于速度每個(gè)體的位置被改變。
第一步是計(jì)算合力F，這將是我們的瓶頸。由于世界上存在的其他物體，單一物體上的力是所有作用力的總和。這導(dǎo)致復(fù)雜度為O（N^2）。速度v和位置r更新的復(fù)雜度都是O（N）。
如果你有興趣，這篇維基百科的文章介紹了一些可以加快力的計(jì)算的近似方法。
純Python
在純Python中，使用NumPy數(shù)組是時(shí)間演變函數(shù)的一種實(shí)現(xiàn)方式，它為優(yōu)化提供了一個(gè)起點(diǎn)，并涉及測(cè)試其他實(shí)現(xiàn)方式。
# 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 |
合力計(jì)算的復(fù)雜度為O（N^2）的現(xiàn)象被NumPy的數(shù)組符號(hào)所掩蓋。每個(gè)數(shù)組操作遍歷數(shù)組元素。
可視化
這里是7個(gè)物體從隨機(jī)初始狀態(tài)開始演化的路徑圖：
性能
為了實(shí)現(xiàn)這個(gè)基準(zhǔn)，我們?cè)陧?xiàng)目目錄下創(chuàng)建了一個(gè)腳本，包含如下內(nèi)容：
import lib
w = lib.World(101)
lib.evolve(w, 4096)
我們使用cProfile模塊來測(cè)試衡量這個(gè)腳本。
python -m cProfile -scum bench.py
前幾行告訴我們，compute_F確實(shí)是我們的瓶頸，它占了超過99％的運(yùn)行時(shí)間。
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臺(tái)式機(jī)上有101體，這種實(shí)現(xiàn)能夠通過每秒257個(gè)時(shí)間步長(zhǎng)演化世界。
簡(jiǎn)單的C擴(kuò)展 1
在本節(jié)中，我們將看到一個(gè)C擴(kuò)展模塊實(shí)現(xiàn)演化的功能。當(dāng)看完這一節(jié)時(shí)，這可能幫助我們獲得一個(gè)C文件的副本。文件src/simple1.c，可以在GitHub上獲得。
關(guān)于NumPy的C API的其他文檔，請(qǐng)參閱NumPy的參考。Python的C API的詳細(xì)文檔在這里。
樣板
文件中的第一件事情是先聲明演化函數(shù)。這將直接用于下面的方法列表。
static PyObject *evolve(PyObject *self, PyObject *args);
接下來是方法列表。
static PyMethodDef methods[] = {
{ "evolve", evolve, METH_VARARGS, "Doc string."},
{ NULL, NULL, 0, NULL } /* Sentinel */
};
這是為擴(kuò)展模塊的一個(gè)導(dǎo)出方法列表。這只有一個(gè)名為evolve方法。
樣板的最后一部分是模塊的初始化。
PyMODINIT_FUNC initsimple1(void) {
(void) Py_InitModule("simple1", methods);
import_array();
}
另外，正如這里顯示，initsimple1中的名稱必須與Py_InitModule中的第一個(gè)參數(shù)匹配。對(duì)每個(gè)使用NumPy API的擴(kuò)展而言，調(diào)用import_array是有必要的。
數(shù)組訪問宏
數(shù)組訪問的宏可以在數(shù)組中被用來正確地索引，無論數(shù)組被如何重塑或分片。這些宏也使用如下的代碼使它們有更高的可讀性。
#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)])
在這里，我們看到訪問宏的一維和二維數(shù)組。具有更高維度的數(shù)組可以以類似的方式被訪問。
在這些宏的幫助下，我們可以使用下面的代碼循環(huán)r:
for(i = 0; i < r_shape(0); ++i) {
for(j = 0; j < r_shape(1); ++j) {
r(i, j) = 0; // Zero all elements.
}
}
命名標(biāo)記
上面定義的宏，只在匹配NumPy的數(shù)組對(duì)象定義了正確的名稱時(shí)才有效。在上面的代碼中，數(shù)組被命名為py_m和py_r。為了在不同的方法中使用相同的宏，NumPy數(shù)組的名稱需要保持一致。
計(jì)算力
特別是與上面五行的Python代碼相比，計(jì)算力數(shù)組的方法顯得頗為繁瑣。
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; } } } |
請(qǐng)注意，我們使用牛頓第三定律（成對(duì)出現(xiàn)的力大小相等且方向相反）來降低內(nèi)環(huán)范圍。不幸的是，它的復(fù)雜度仍然為O（N^2）。
演化函數(shù)
該文件中的最后一個(gè)函數(shù)是導(dǎo)出的演化方法。
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參數(shù)如何被解析。在該函數(shù)底部的時(shí)間步長(zhǎng)循環(huán)中，我們看到的速度和位置向量的x和y分量的顯式計(jì)算。
性能
C版本的演化方法比Python版本更快，這應(yīng)該不足為奇。在上面提到的相同的i5臺(tái)式機(jī)中，C實(shí)現(xiàn)的演化方法能夠?qū)崿F(xiàn)每秒17972個(gè)時(shí)間步長(zhǎng)。相比Python實(shí)現(xiàn)，這方面有70倍的提升。
觀察
注意，C代碼一直保持盡可能的簡(jiǎn)單。輸入?yún)?shù)和輸出矩陣可以進(jìn)行類型檢查，并分配一個(gè)Python裝飾器函數(shù)。刪除分配，不僅能加快處理，而且消除了由Python對(duì)象不正確的引用計(jì)數(shù)造成的內(nèi)存泄露（或更糟）。
下一部分
在本系列文章的下一部分，我們將通過發(fā)揮C-相鄰NumPy矩陣的優(yōu)勢(shì)來提升這種實(shí)現(xiàn)的性能。之后，我們來看看使用英特爾的SIMD指令和OpenMP來進(jìn)一步推進(jìn)。