numpy实现等步长4阶Runge Kutta算法用来求解SEIR模型
1. Runge Kutta数值方法
《数值计算方法》（第2版）清华大学出版社[1]给出了一个未知数的4阶Runge Kutta公式：
{
y
n
+
1
=
y
n
+
h
6
(
k
1
+
2
k
2
+
2
k
3
+
k
4
)
,
k
1
=
h
f
(
x
n
,
y
n
)
,
k
2
=
h
f
(
x
n
+
h
2
,
y
n
+
k
1
2
)
,
k
3
=
h
f
(
x
n
+
h
2
,
y
n
+
k
2
2
)
,
k
4
=
h
f
(
x
n
+
h
2
,
y
n
+
k
3
2
)
,
\begin{cases} y_{n+1} = y_n + \frac{h}{6}(k_1+2k_2+2k_3+k_4),\\k_1=hf(x_n,y_n),\\ k_2=hf(x_n+\frac{h}{2},y_n+\frac{k_1}{2}),\\ k_3=hf(x_n+\frac{h}{2},y_n+\frac{k_2}{2}),\\ k_4=hf(x_n+\frac{h}{2},y_n+\frac{k_3}{2}), \end{cases}
⎩⎪⎪⎪⎪⎪⎪⎨⎪⎪⎪⎪⎪⎪⎧yn+1=yn+6h(k1+2k2+2k3+k4),k1=hf(xn,yn),k2=hf(xn+2h,yn+2k1),k3=hf(xn+2h,yn+2k2),k4=hf(xn+2h,yn+2k3),
对于两个未知数的微分方程组
{
d
y
d
x
=
f
(
x
,
y
,
z
)
,
y
(
x
0
)
=
y
0
,
d
z
d
x
=
g
(
x
,
y
,
z
)
,
z
(
x
0
)
=
z
0
\begin{cases} \frac{\mathrm dy}{\mathrm dx}=f(x,y,z),\quad y(x_0)=y_0,\\ \frac{\mathrm dz}{\mathrm dx}=g(x,y,z),\quad z(x_0)=z_0\\ \end{cases}
{dxdy=f(x,y,z),y(x0)=y0,dxdz=g(x,y,z),z(x0)=z0
Runge Kutta公式变为：
{
y
n
+
1
=
y
n
+
h
6
(
k
1
+
2
k
2
+
2
k
3
+
k
4
)
,
z
n
+
1
=
z
n
+
h
6
(
m
1
+
2
m
2
+
2
m
3
+
m
4
)
,
k
1
=
h
f
(
x
n
,
y
n
,
z
n
)
,
k
2
=
h
f
(
x
n
+
h
2
,
y
n
+
k
1
2
,
z
n
+
m
1
2
)
,
k
3
=
h
f
(
x
n
+
h
2
,
y
n
+
k
2
2
,
z
n
+
m
2
2
)
,
k
4
=
h
f
(
x
n
+
h
,
y
n
+
k
3
,
z
n
+
m
3
)
,
m
1
=
h
g
(
x
n
,
y
n
,
z
n
)
,
m
2
=
h
g
(
x
n
+
h
2
,
y
n
+
k
1
2
,
z
n
+
m
1
2
)
,
m
3
=
h
g
(
x
n
+
h
2
,
y
n
+
k
2
2
,
z
n
+
m
2
2
)
,
m
4
=
h
g
(
x
n
+
h
,
y
n
+
k
3
,
z
n
+
m
3
)
,
n
=
0
,
1
,
2
,
⋯
\begin{cases} y_{n+1} = y_n + \frac{h}{6}(k_1+2k_2+2k_3+k_4),\\ z_{n+1} = z_n + \frac{h}{6}(m_1+2m_2+2m_3+m_4),\\ k_1=hf(x_n,y_n,z_n),\\ k_2=hf(x_n+\frac{h}{2},y_n+\frac{k_1}{2},z_n+\frac{m_1}{2}),\\ k_3=hf(x_n+\frac{h}{2},y_n+\frac{k_2}{2},z_n+\frac{m_2}{2}),\\ k_4=hf(x_n+h,y_n+k_3,z_n+m_3),\\ m_1=hg(x_n,y_n,z_n),\\ m_2=hg(x_n+\frac{h}{2},y_n+\frac{k_1}{2},z_n+\frac{m_1}{2}),\\ m_3=hg(x_n+\frac{h}{2},y_n+\frac{k_2}{2},z_n+\frac{m_2}{2}),\\ m_4=hg(x_n+h,y_n+k_3, z_n+m_3), \end{cases} n=0,1,2,\cdots
⎩⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎨⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎧yn+1=yn+6h(k1+2k2+2k3+k4),zn+1=zn+6h(m1+2m2+2m3+m4),k1=hf(xn,yn,zn),k2=hf(xn+2h,yn+2k1,zn+2m1),k3=hf(xn+2h,yn+2k2,zn+2m2),k4=hf(xn+h,yn+k3,zn+m3),m1=hg(xn,yn,zn),m2=hg(xn+2h,yn+2k1,zn+2m1),m3=hg(xn+2h,yn+2k2,zn+2m2),m4=hg(xn+h,yn+k3,zn+m3),n=0,1,2,⋯
1.1 MATLAB实现
以下是4阶Runge Kutta算法的MATLAB实现（对《数值计算方法》322页给出的MATLAB脚本略微修改得到）：
function [T, Z] = rks4(F, a, b, Za, M)
% f 是目标函数，[a b] 是左右端点， Za = [f1(a), ..., fn(a)] 是初始值
% M 是步长（步数），T 是步长矢量，Z = [f1(t), ..., fn(t)] 是近似值
h = (b -a)/ M;
T = zeros(1, M+1);
Z = zeros(M+1, length(Za));
T = a: h: b;
Z(1,:) = Za;
for i = 1: M
k1 = h* feval(F, T(i), Z(i,:))';
k2 = h* feval(F, T(i)+h/2, Z(i,:)+k1/2)';
k3 = h* feval(F, T(i)+h/2, Z(i,:)+k2/2)';
k4 = h* feval(F, T(i)+h, Z(i,:)+k3)';
Z(i+1,:) = Z(i,:) + (k1 + 2*k2 + 2*k3 + k4)/6;
end
% T要和Z的维度一致，返回结果按行为每一步迭代结果、按列为各函数取值
T = T';
end
1.2 python实现
把上面MATLAB脚本用numpy实现，编写成自己的包并在main()函数测试求解微分方程组
{
d
x
d
t
=
x
,
x
(
0
)
=
1
,
d
y
d
t
=
y
,
y
(
0
)
=
1
\begin{cases} \frac{\mathrm dx}{\mathrm dt}=x,\quad x(0)=1,\\ \frac{\mathrm dy}{\mathrm dt}=y,\quad y(0)=1\\ \end{cases}
{dtdx=x,x(0)=1,dtdy=y,y(0)=1
他的解为
x
=
y
=
e
t
x=y=e^t
x=y=et
# -*-coding: utf-8-*-
# 四阶Runge Kutta方法
__all__ = ['rks4']
import numpy as np
# 包含在'''''' 或 """""" 之间且不属于任何语句的内容将被解释器认为是注释
"""
Author: Guo
Step identical 等步长 Runge Kutta 4
Usage: tval, yval = rks4(dYdt, t_start, t_end, Y_t0, steps -1)
Parameters:
f -- 要求函数f返回导数值的列向量（与MATLAB风格一致）
Za -- 初始条件的numpy array. (np.array([f1_0, f2_0, f3_0, ...]))
Return:
t -- A column vector of t values at each step
Z -- A matrix containing column vectors of all function values at each step
"""
def rks4(f, a, b, Za, M):
# M + 1 steps in total
h = (b - a)/ M
t = np.linspace(a, b, M + 1).reshape(M+1,1)
Z = np.zeros((M+1, Za.size))
Z[0,:] = Za
# print(Z.shape,Z[:,0].shape, Z[0,:].shape)
for i in range(1,M+1):
k1 = h * f(t[i-1], np.transpose(Z[i-1,:]))
k2 = h * f(t[i-1] +h/2, np.transpose(Z[i-1,:] +k1/2))
k3 = h * f(t[i-1] +h/2, np.transpose(Z[i-1,:] +k2/2))
k4 = h * f(t[i-1] +h, np.transpose(Z[i-1,:] +k3))
# print(k4.shape)
Z[i,:] = Z[i-1,:] + (k1 + 2*k2 + 2*k3 + k4)/6
# print(t.shape, Z.shape, sep='\n')
return (t, Z)
def main():
# 微分方程 dxdt = x; dydt = y
# 返回dxdt、dydt的列向量
ode = lambda t, y: np.array([y[0], y[1]]).T
tval, yval = rks4(ode, 0 ,1, np.array([1,1]) ,5)
# np.diff(ans[::1000,0])
y = np.exp(np.linspace(0,1,yval.shape[0]))
# print(yval.shape[0] == y.shape[0])
print(np.hstack((tval, yval)))
s = np.std(yval[:,0] - y)/np.mean(y) *100
print(f"Raletive Error: {s}%")
if __name__ == '__main__':
# main()
# 运行main()函数 请注释掉此语句
exit('Please use rks4 as module')
[[0. 1. 1. ]
[0.2 1.2214 1.2214 ]
[0.4 1.49181796 1.49181796]
[0.6 1.82210646 1.82210646]
[0.8 2.22552083 2.22552083]
[1. 2.71825114 2.71825114]]
Raletive Error: 0.00060603304842948%
4阶Runge Kutta公式的截断误差为 O ( h 5 ) O(h^5) O(h5)，可以看到结果与书上介绍基本一致。
2. 用scipy.optimize least_sqares方法拟合SEIR模型参数并求解
2.1 scipy.optimize.least_squares(fun, x0）
-
功能：给定计算残差的函数和超参数初始值，用最小二乘方法给出超参数的最优值（局部）。
-
fun:callable
fun应当具备这样的基本形式： f u n ( x , ∗ a r g s , ∗ ∗ k w a r g s ) fun(x, *args, **kwargs) fun(x,∗args,∗∗kwargs)（就是说 l e a s t _ s q u a r e least\_square least_square函数的最小化过程对fun的第一个参数x进行），x必须是np.array类型。返回值必须也是一维数组np.array。
官方文档还解释了fun的返回值为复数的处理方法： If the argument x is complex or the function fun returns complex residuals, it must be wrapped in a real function of real arguments, as shown at the end of the Examples section. -
x0: np.array类型。给出参数初始值。
-
返回值：
| 变量 |
含义
|
|---|---|
| status: int | The reason for algorithm termination: -1 : improper input parameters status returned from MINPACK. 0 : the maximum number of function evaluations is exceeded. 1 : gtol termination condition is satisfied. 2 : ftol termination condition is satisfied. 3 : xtol termination condition is satisfied. 4 : Both ftol and xtol termination conditions are satisfied. |
| message: str | Verbal description of the termination reason. |
| success: bool | True if one of the convergence criteria is satisfied (status > 0) |
2.2 拟合SEIR模型参数的python脚本
# -*-coding: utf-8-*-
# 使用scipy.optimize.least_squares进行最小二乘法拟合参数
import sys
"""对于自己编写的模块可以通过sys.path.append()把对应路径加入python搜索模块路径"""
sys.path.append(r"package mcmutils所在路径名")
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import seaborn as sns
from mcmutils import odes
data = pd.read_csv(r"./china.csv")
# data.head()
N = 140005e4
I = np.array(data.existing)
R = np.array(data.recovered[:-6])
E = np.array(I[6:] - I[:-6])
I = I[:-6]
S = N*np.ones_like(E) - E - I - R
y = np.array([S, E, I, R])
# 注意numpy的索引是多维数组风格
st = np.argmax(y[2,:])
# st = np.argmax(np.diff(y[2,:]))
# # st = 0
y = np.transpose(y[:,st:])
def seir(t, y, arg, N):
(beta, lamda, sigma, mu, alpha) = (arg[0], arg[1], arg[2], arg[3], arg[4])
dS = -beta*lamda*(y[1]+y[2])*y[0]/N
dE = beta*lamda*(y[1]+y[2])*y[0]/N - y[1]/sigma
dI = y[1]/sigma - (mu+alpha)*y[2]
dR = (mu+alpha)*y[2]
# 按照MATLAB风格 返回导数值的列向量
return np.array([dS, dE, dI, dR]).T
# 残差函数
def rfun(args, ydata, N):
len = ydata.shape[0]
y0 = ydata[0,:]
# print(y0, len)
_, yval = odes.rks4(lambda t,y: seir(t,y, args, N), 1, len, y0, len-1)
# print(yval.shape)
# return np.log(np.sum((yval[:,1:] - ydata)**2))
weights = np.array([0.1,0.1,0.6,0.4])
return np.sum(np.abs(yval - ydata)*weights, axis=1)
# print(y.shape)
# rfun(np.random.rand(5,1), y)
from scipy.optimize import least_squares
arg0 = np.array([0.1, 3, 7, 0.1, 1e-3])
optiargs = least_squares(lambda arg: rfun(arg, y, N), arg0)
print(optiargs)
len = y.shape[0]
y0 = y[0,:]
# print(y0)
tval, yval = odes.rks4(lambda t,y: seir(t,y, optiargs.x, N), 1, len, y0, len -1)
sns.set()
fig, ax = plt.subplots()
plt.plot(tval, yval[:,2], 'r--', linewidth=1.5, label='Fit')
plt.plot(tval, y[:,2], 'k--', linewidth=1.5, label='True')
plt.legend(loc='best')
plt.xlabel('days')
plt.ylabel('cases')
plt.title('SEIR model fitting')
err = np.std(yval[:,2]- y[:,2])/np.mean(y[:,2])*100
ax.text(0.4, 0.9, f'Relative error:\n{err:.{3}}%', transform=ax.transAxes)
plt.show()
2.3 求解结果
message: '`xtol` termination condition is satisfied.'
nfev: 66
njev: 43
optimality: 261400320.2678833
status: 3
success: True
x: array([0.11668877, 3.00386719, 2.3342517 , 1.39832263, 1.2959752 ])
参考资料
[1]. 《数值计算方法》（第2版）清华大学出版社293-322
3346
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
