建立传染病SIR模型代码

import scipy.integrate as spi
import numpy as np
import matplotlib.pyplot as plt

beta = 7e-3
gamma = 4e-3
TS = 1.0
ND = 1000.0
TN = 4.2
S0 = 0.998*TN
I0 = 0.002*TN
INPUT = (S0, I0, 0.0)


def diff_eqs(INP, t):
    Y = np.zeros((3))
    V = INP
    Y[0] = - beta * V[0] * V[1]
    Y[1] = beta * V[0] * V[1] - gamma * V[1]
    Y[2] = gamma * V[1]
    return Y


t_start = 0.0
t_end = ND
t_inc = TS
t_range = np.arange(t_start, t_end + t_inc, t_inc)
RES = spi.odeint(diff_eqs, INPUT, t_range)

print(RES)
with open('SIR.txt', 'w') as f:
    for each in RES:
        f.write(str(each))
        f.write('\n')
f.close()
# Ploting
plt.subplot(111)
plt.plot(RES[:, 1], '-r', label='Infectious')
plt.plot(RES[:, 0], '-g', label='Susceptibles')
plt.plot(RES[:, 2], '-k', label='Recovereds')
plt.legend(loc=0)
plt.title('SIR')
plt.xlabel('Time (day)')
plt.ylabel('Infectious Susceptibles and Recovereds (million)')
plt.show()

其中$\beta$为感染率，$\gamma$为药物有效性，TS和ND分别为时间间隔和结束时间，TN为区域内人口数（单位为百万）。
运行结果为