如何在Python中实现多组参数的ODE方程组求解（Matlab转译）

问题背景 #

刚接触Python，正在将Matlab代码转译为Python，需要用多组不同参数运行僵尸模型，参数如下：

P = [0, 0.5, 0.7];
d = [0.001, 0.002, 0.003];
B = [0.0095, 0.0080, 0.0070];
G = [0.001, 0.002, 0.003];
A = [0.001, 0.002, 0.003];

在Matlab中可通过嵌套函数结合for循环轻松完成三次不同参数的模型运行，但在Python中未找到类似实现的相关文档，需修改提供的Python僵尸ODE代码，达到与Matlab相同的运行效果。原Python代码如下：

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
plt.ion()
plt.rcParams['figure.figsize'] = 10, 8

P = 0      # birth rate
d = 0.0001  # natural death percent (per day)
B = 0.0095  # transmission percent  (per day)
G = 0.0001  # resurect percent (per day)
A = 0.0001  # destroy percent  (per day)

# solve the system dy/dt = f(y, t)
def f(y, t):
     Si = y[0]
     Zi = y[1]
     Ri = y[2]
     # the model equations (see Munz et al. 2009)
     f0 = P - B*Si*Zi - d*Si
     f1 = B*Si*Zi + G*Ri - A*Si*Zi
     f2 = d*Si + A*Si*Zi - G*Ri
     return [f0, f1, f2]

# initial conditions
S0 = 500.              # initial population
Z0 = 0                 # initial zombie population
R0 = 0                 # initial death population
y0 = [S0, Z0, R0]     # initial condition vector
t  = np.linspace(0, 5., 1000)         # time grid

# solve the DEs
soln = odeint(f, y0, t)
S = soln[:, 0]
Z = soln[:, 1]
R = soln[:, 2]

解决思路 #

Python中可通过将参数作为额外参数传入ODE函数，配合for循环遍历参数组的方式，替代Matlab的嵌套函数实现，完成多组参数的模型运行。

修改后的代码 #

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
plt.ion()
plt.rcParams['figure.figsize'] = 10, 8

# 定义多组参数，按索引对应每组的参数值
params = {
    'P': [0, 0.5, 0.7],
    'd': [0.001, 0.002, 0.003],
    'B': [0.0095, 0.0080, 0.0070],
    'G': [0.001, 0.002, 0.003],
    'A': [0.001, 0.002, 0.003]
}

# 修改ODE函数，加入参数作为输入
def f(y, t, P, d, B, G, A):
    Si = y[0]
    Zi = y[1]
    Ri = y[2]
    # 模型方程
    f0 = P - B*Si*Zi - d*Si
    f1 = B*Si*Zi + G*Ri - A*Si*Zi
    f2 = d*Si + A*Si*Zi - G*Ri
    return [f0, f1, f2]

# 初始条件
S0 = 500.
Z0 = 0
R0 = 0
y0 = [S0, Z0, R0]
t = np.linspace(0, 5., 1000)

# 遍历每组参数运行模型
for i in range(len(params['P'])):
    # 获取当前组的参数
    current_P = params['P'][i]
    current_d = params['d'][i]
    current_B = params['B'][i]
    current_G = params['G'][i]
    current_A = params['A'][i]
    
    # 传入参数求解ODE，用args传递额外参数
    soln = odeint(f, y0, t, args=(current_P, current_d, current_B, current_G, current_A))
    S = soln[:, 0]
    Z = soln[:, 1]
    R = soln[:, 2]
    
    # 绘制当前组的结果（可选）
    plt.plot(t, S, label=f'组{i+1}: 人类')
    plt.plot(t, Z, label=f'组{i+1}: 僵尸')
    plt.plot(t, R, label=f'组{i+1}: 死亡')

plt.xlabel('时间（天）')
plt.ylabel('数量')
plt.title('多组参数下的僵尸启示录模型')
plt.legend()
plt.show()

关键说明 #

• 参数传递：利用odeint的args参数，将每组参数传递给ODE函数f，避免依赖全局变量，提升函数灵活性。
• 循环遍历：通过for循环遍历参数组索引，依次取出对应参数值，运行模型并处理结果（如绘图）。
• 扩展性：若需添加更多参数组，只需在params字典中补充对应值即可。

内容的提问来源于stack exchange，提问作者Tee