Python generator 实现 Gillespie 算法

Gillespie 算法与连续时间 Markov 过程

Gillespie 算法是一种可以用于采样连续时间 Markov 过程（CTMC）的随机轨迹的算法。这一算法被广泛用于模拟化学反应网络、生物物理过程等随机过程。在介观热力学中，Gillespie 算法被广泛应用于模拟主方程描述的系统，即连续时间 Markov 系统。
对于连续时间、离散状态的马尔可夫过程，其概率演化可以由如下主方程表达

$$ \frac{\mathrm d p(x, t)}{\mathrm dt} = R p(x, t). $$

其中矩阵 $R$ 的矩阵元满足

$$ \begin{cases} R_{ij} \ge 0, \quad i \neq j; \\ \\ R_{ii} = \displaystyle{\sum_{j\neq i}}R_{ij}. \end{cases} $$

列和为零的性质保证的概率的守恒（即归一性）。
上式表明我们可以将连续时间 Markov 过程视为概率密度在各个状态上的流动过程。实际上，我们也可以考虑大量按初始概率分布的无相互作用粒子，粒子在不同状态之间的转移速率由转移矩阵 $R$ 给出。每个粒子在所有状态上的运动都构成一条随机轨迹，每一时刻所有粒子的分布情况即为该时刻系统概率密度的演化情况。
Gillespie 算法即是基于上述的第二种视角，该算法可以模拟单个粒子在系统各个状态上的随机轨迹，由此得到的随机轨迹是上述主方程的解。随机轨迹本身给出系统的动力学性质，重复对轨迹进行采样，又可由轨迹平均得到系统的一系列统计性质。因此，Gillespie 算法之于介观热力学的重要性不言而喻。

算法原理

对于主方程描述的系统，粒子在某一状态上向其相邻状态的转移概率于转移矩阵 $R$ 有关。一般地，定义矩阵元 $R_{ij}$ 表示状态 $j$ 向状态 $i$ 的转移速率。假设在极短的时间内，粒子仅可能发生一步跳跃，即发生多步跳跃的概率为时间的高阶无穷小量。这样的概率模型实际上是 Poisson 过程，即粒子每一步的跳跃过程都是一个 Poisson 过程。从状态 $j$ 到状态 $i$ 的 Poisson 过程的等待时间满足 Poisson 分布

$$ P_{i\leftarrow j}(T \le t) = 1 - e^{-R_{ij}t}. $$

对于一般的 Markov 系统，与某一状态 $i$ 相邻的状态可能有很多，粒子跳出状态 $i$ 的概率应当为跳到任意与 $i$ 相邻的状态的概率之和。由于 Poisson 分布之和依然是 Poisson 分布，因此总的跳出过程依然是 Poisson 过程，且其速率为所有可能的跳出速率之和。由矩阵 $R$ 的矩阵元定义可知，所有从状态 $i$ 跳出的速率之和即为 $-R_{ii}$。
由上述分析可知，粒子在状态 $i$ 的等待时间满足分布

$$ P(T \le t) = 1 - e^{R_{ii}t}. $$

在等待时间之后，粒子跳到不同状态的概率与各个转移速率成正比。待粒子到了新的状态，再重复上述过程，即可得到一条随机轨迹。

算法细节

每一步跳跃都涉及到两个随机的因素，即随机的等待时间和随机的跳跃方向。我们可以通过生成两个随机数来实现这两个过程。
粒子在状态 $i$ 上的等待时间可以由一个在 $[0, 1)$ 上均匀分布的随机数 $r_1$ 生成。即

$$ \text{waiting time} = \frac{1}{R_{ii}}\ln r_1. $$

粒子下一步所处的位置可以由另一个在 $[0, 1)$ 上均匀分布的随机数 $r_2$ 生成。粒子下一步所处的位置为满足如下表达式的位置 $j'$

$$ \frac{1}{-R_{ii}}\sum_j^{j'}R_{ji} < r <\frac{1}{-R_{ii}}\sum_j^{j'+1}R_{ji}. $$

这表明下一步处在某一状态的概率与其对应的转移速率成正比。
算法有如下伪代码

Step 1. Initialize t = t0, x = x0
Step 2. Calculate R_ji and R_ii
Step 3. Generate waiting time with R_ii and r_1
Step 4. Generate next state with R_ii, R_ji and r_2
Step 5. Record results, return to step 1 or end the loop

Python 生成器实现算法

生成器是一种特殊的 Python 结构，它可以当作普通的可迭代对象使用，比如用在循环中。相比于一个普通的循环函数而言，生成器更加的节省内存空间，代码更加简洁优美，且具有更好的泛用性和拓展性。同时，Numba 等即时编译器也支持生成器的使用，为代码运行速度提供了保障。综上所述，我们选用生成器来实现 Gillespie 算法。
对于任意矩阵 $R$ 和初始概率分布 $p_{\text{ini}}$，如下算法构建了一个随机轨迹生成器。

import numpy as np
import numba as nb

@nb.jit(nopython=True)
def traj(R, p_ini):
    
    # total number of states
    n = len(R[0, :])
    
    # modify R with diagonal elements 0
    for i in range(n):
        R[i, i] = 0
    
    # choose a state based on the initial prob distribution
    r_1 = np.random.rand()
    state = np.searchsorted(np.cumsum(p_ini), r_1, side="right")
    
    # calculate waiting time with a random number
    a_i = R[:, state]
    sum_a_i = np.sum(a_i)
    r_2 = np.random.rand()
    waiting_time = -np.log(r_2) / sum_a_i
    
    # yield current state and waiting time
    yield state, waiting_time
    
    # start the loop
    while True:
        # generate the next state
        r_1 = np.random.rand()
        state = np.searchsorted(np.cumsum(a_i/sum_a_i), r_1, side="right")
        
        # calculate the waiting time on the next state
        a_i = R[:, state]
        sum_a_i = np.sum(a_i)
        r_2 = np.random.rand()
        waiting_time = -np.log(r_2) / sum_a_i
        
        # yield next state and waiting time
        yield state, waiting_time

上述代码又两部分组成。第一部分在 while True 之前，这部分由初始概率分布生成处在某一位置的粒子，并计算粒子在该状态上的等待时间。第二部分在 while True 之后，由一个循环构成，每次使用 next() 函数调用该生成器都将进行一步迭代，生成一次跳跃过程及其对应的等待时间。函数开头的修饰器 @nb.jit(nopython=True) 调用了即时编译器 Numba，它通过即使编译该函数的方式大幅提高其运行速度，代价是函数内容需要仔细编写以满足 Numba 类型推断器的要求。
可以使用 Python 内置的 next() 函数调用该生成器。例如可以用如下代码生成一段时间在 $[0, 100]$ 的随机轨迹

T = 100
x = []
t = []
current_time = 0
trajectory = traj(R, p_ini)
while current_time < T:
    state, waiting_time = next(trajectory)
    x.append(state)
    current_time += waiting_time
    t.append(current_time)