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

Posted 2023-03-12

【中文标题】通过 Python 并行求解具有大量初始条件的 ODE

【英文标题】：Solving ODE with large set of initial conditions in parallel through Python

【发布时间】：2020-01-02 12:56:42

【问题描述】：

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

【问题讨论】：

为什么二维初始条件会有所改进？ ode 函数是否尽可能快？这是您对评估速度的最大控制权。

对于大规模生产来说，scipy 代码不是最理想的。使用绑定到日晷或DifferentialEquations.jl（有一些重叠）来访问最新的和经过测试的实现。对 ODE 函数使用 JIT 或直接编译语言。请注意，并行解决方案可能会影响单个解决方案的准确性，尤其是。如果解集合访问了状态空间的不同部分。

您不需要二维数组函数：您可以使用一维数组输入来制定 ODE 函数，然后在函数代码的第一行对其进行整形。正如 hpaulj 所说，您可以通过使用 ode 函数来提高速度。您可以尝试的其他事情是：减少收敛标准或使用显式积分器，如 RK23 或 RK45 以最小步长，通过雅可比

“并行计算”位见***.com/questions/4682429/parfor-for-python（假设不涉及跨线程通信）。

不如可以接受多个初始值的单个函数调用方便，但可以选择多处理：***.com/questions/34291639/…

【参考方案1】：

以下是使用 Python 并行求解具有大量初始条件的 ODE 的两个示例。首先，最快的是使用带有NumbaLSODA ODE 集成器的 Numba 多线程。这里我使用 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()