用parareal算法求解常微分方程的python

以下是使用 parareal 算法求解常微分方程的 Python 代码示例：

import numpy as np
from scipy.integrate import solve_ivp
from joblib import Parallel, delayed

def f(t, y):
    return np.array([y[1], -y[0]])

# sequential solver
def solver_seq(f, tspan, y0, h):
    sol = solve_ivp(f, tspan, y0, t_eval=np.arange(tspan[0], tspan[1], h))
    return sol.y.T

# coarse solver
def solver_coarse(f, tspan, y0, h_coarse):
    sol = solve_ivp(f, tspan, y0, t_eval=np.arange(tspan[0], tspan[1], h_coarse))
    return sol.y.T

# fine solver
def solver_fine(f, tspan, y0, h_fine):
    sol = solve_ivp(f, tspan, y0, t_eval=np.arange(tspan[0], tspan[1], h_fine))
    return sol.y.T

# parareal solver
def solver_parareal(f, tspan, y0, h_coarse, h_fine, max_iter=10, tol=1e-6, n_jobs=-1):
    y_coarse = solver_coarse(f, tspan, y0, h_coarse)
    y_fine = solver_fine(f, tspan, y0, h_fine)
    error = np.linalg.norm(y_fine - y_coarse, np.inf)
    iteration = 0
    while error > tol and iteration < max_iter:
        y_parareal = Parallel(n_jobs=n_jobs)(
            delayed(solve_ivp)(f, [tspan[i], tspan[i+1]], y_coarse[i], t_eval=[tspan[i+1]], method='RK45').y[:, 0]
            for i in range(len(tspan)-1))
        y_parareal = np.vstack((y0, np.array(y_parareal)))
        y_coarse = y_coarse + (y_fine - y_parareal)
        y_fine = solver_fine(f, tspan, y0, h_fine)
        error = np.linalg.norm(y_fine - y_coarse, np.inf)
        iteration += 1
    return y_coarse

# example usage
tspan = [0, 10]
y0 = [0, 1]
h_coarse = 0.5
h_fine = 0.01
y_seq = solver_seq(f, tspan, y0, h_fine)
y_parareal = solver_parareal(f, tspan, y0, h_coarse, h_fine, n_jobs=4)
print(y_seq)
print(y_parareal)

其中，`f` 函数定义了常微分方程的右侧函数，`solver_seq` 是顺序求解器，`solver_coarse` 和 `solver_fine` 分别是粗网格和细网格求解器，`solver_parareal` 是 Parareal 求解器。在 Parareal 求解器中，我们首先用粗网格求解器求解一次，然后再用细网格求解器求解一次，将两次求解的结果之差作为 Parareal 残差，然后使用多进程运行 Parareal 迭代，并且在每个时间步上使用 RK45 算法求解 Parareal 残差。最后返回 Parareal 求解器的解 `y_coarse`。

在例子中，我们使用了一个简单的常微分方程 `y'' + y = 0`，其解为 `y = sin(t)`。我们使用 `solver_seq` 求解器求解精确解，并使用 `solver_parareal` 求解器求解 Parareal 解。可以看到，两个解的结果非常接近。