通过 Python 并行求解具有大量初始条件的 ODE

Question

我正在尝试使用scipy.integrate.ode或scipy.integrate.odeint求解大量（超过一千个）初始条件的 ODE 系统，但是执行循环的速度非常慢，并且 scipy 似乎没有提供用于输入 2D 数组（由一组指定初始条件的一维数组），并且选项vectorized似乎scipy.integrate.solve_ivp并不意味着它接受初始条件的二维数组（https://docs.scipy.org/doc/scipy/reference/ generated/scipy.integrate .solve_ivp.html）。我读过一个线程询问类似的问题（Vectorized SciPy odesolver），其中一个答案建议使用scipy.integrate.odeint，但是它似乎也不接受多维数组，所以它根本不明白如何实现它。有什么解决方案可以加快进程吗？除了矢量化之外，我还考虑过使用并行计算技术，但我对此并不熟悉，而且我认为它并没有真正像矢量化那样显着地加速程序？

Answer 1

以下是使用 Python 并行求解具有大量初始条件的 ODE 的两个示例。首先，最快的方法是将 Numba 多线程与NumbaLSODA ODE 积分器结合使用。这里我使用 Numba 来编译所有代码，因此循环速度非常快。

此示例需要 0.175 秒。

from NumbaLSODA import lsoda_sig, lsoda
from matplotlib import pyplot as plt
import numpy as np
import numba as nb

@nb.cfunc(lsoda_sig)
def f(t, u, du, p):
    du[0] = u[0]-u[0]*u[1]
    du[1] = u[0]*u[1]-u[1]

funcptr = f.address
t_eval = np.linspace(0.0,20.0,201)
np.random.seed(0)

@nb.njit(parallel=True)
def main(n):
    u1 = np.empty((n,len(t_eval)), np.float64)
    u2 = np.empty((n,len(t_eval)), np.float64)
    for i in nb.prange(n):
        u0 = np.empty((2,), np.float64)
        u0[0] = np.random.uniform(4.5,5.5)
        u0[1] = np.random.uniform(0.7,0.9)
        usol, success = lsoda(funcptr, u0, t_eval, rtol = 1e-8, atol = 1e-8)
        u1[i] = usol[:,0]
        u2[i] = usol[:,1]
    return u1, u2

u1, u2 = main(10000)

plt.rcParams.update({'font.size': 15})
fig,ax = plt.subplots(1,1,figsize=[7,5])
low, med, high = np.quantile(u1,(.025,.5,.975),axis=0)
ax.plot(t_eval,med)
ax.fill_between(t_eval,low,high,alpha=0.3)
low, med, high = np.quantile(u2,(.025,.5,.975),axis=0)
ax.plot(t_eval,med)
ax.fill_between(t_eval,low,high,alpha=0.3)
plt.show()

多处理和scipy.integrate.odeint. 此示例需要 2.9 秒（比 NumbaLSODA 慢 16 倍）。

from scipy.integrate import odeint
import numba as nb
import numpy as np
from matplotlib import pyplot as plt
from pathos.multiprocessing import ProcessingPool as Pool

@nb.njit
def f_sp(u, t):
    return np.array([u[0]-u[0]*u[1],u[0]*u[1]-u[1]])

t_eval = np.linspace(0.0,20.0,201)

def main(u0):
    usol = odeint(f_sp, u0, t_eval, rtol = 1e-8, atol = 1e-8)
    return usol[:,0], usol[:,1]

n = 10000
u0_1 = np.random.uniform(4.5,5.5,n).reshape((n,1))
u0_2 = np.random.uniform(0.7,0.9,n).reshape((n,1))
u0_all = np.append(u0_1, u0_2, axis=1)

p = Pool(6)
sol = p.map(main, u0_all)

u1 = np.empty((n,len(t_eval)), np.float64)
u2 = np.empty((n,len(t_eval)), np.float64)
for i in range(n):
    u1[i] = sol[i][0]
    u2[i] = sol[i][1]

plt.rcParams.update({'font.size': 15})
fig,ax = plt.subplots(1,1,figsize=[7,5])
low, med, high = np.quantile(u1,(.025,.5,.975),axis=0)
ax.plot(t_eval,med)
ax.fill_between(t_eval,low,high,alpha=0.3)
low, med, high = np.quantile(u2,(.025,.5,.975),axis=0)
ax.plot(t_eval,med)
ax.fill_between(t_eval,low,high,alpha=0.3)
plt.show()