Python 计算思维训练——常微分新冠预测

第2关：新冠病毒传播预测-隔离与不隔离

任务描述

本关任务：使用常微分方程来预测新冠疫情，加上隔离与不隔离的情况区分。

相关知识

本关知识点见上一关卡。

编程要求

根据提示，在右侧编辑器 Begin-End 区间补充代码，实现新冠疫情的预测。

测试说明

平台会对你编写的代码进行测试：

---

开始你的任务吧，祝你成功！

import ode_7 as ode
import numpy as np
import matplotlib.pyplot as plt
# 采用matplotlib作图时默认设置下是无法显示中文的，凡是汉字都会显示成方块。
# 实际上，matplotlib是支持unicode编码的，不能正常显示汉字主要是没有找到合适的中文字体。
from pylab import mpl
mpl.rcParams['font.sans-serif']= ['SimHei']
# 解决负号显示问题
import matplotlib
matplotlib.rcParams['axes.unicode_minus']=False

'''
根据传染病传播的SIR模型，微分方程组可以转换为向量运算：
y = [S, I, R]
f(t, y) = [-beta*y[0]*y[1], beta*y[0]*y[1]-gamma*y[1], gamma*y[1]]
前向欧拉公式：
    y[k+1] = y[k] + dt * f(t[k], y[k])
后向欧拉公式（两步迭代）：
    y_p = y[k] + dt * f(t[k], y[k])
    y[k+1] = y[k] + dt * f(x[k+1], y_p)
改进的欧拉法：
    y_p = y[k] + dt * f(t[k], y[k])
    y_c = y[k] + dt * f(t[k+1], y_p)
    y[k+1] = 1/2 * (y_p + y_c)
'''
class SIR(ode.ODE):
    def f(self, t, y): # Y为向量[S,I,R]，t为标量
        return np.array([-self.beta*y[0]*y[1], self.beta*y[0]*y[1]-self.gamma*y[1], self.gamma*y[1]])

    def __init__(self, beta, gamma, y0):
        self.beta = beta; self.gamma = gamma; self.y0 = y0  # 参数属性
        ode.ODE.__init__(self, self.f, self.y0)

N     = 1e8             # 武汉总人数：1000万人
beta  = 1.0/N           # 1名病毒携带者平均每天感染1.0人
y0    = [N-1, 1, 0]     # 初始发病1人，其他人员正常 [S0, I0, R0]
t = np.arange(0, 61, 1) # 模拟60天的发展情况，单位时间为1天

# 对不同隔离机制进行对比分析：

# 未实施隔离： gamma = 1/25，假设肺炎平均25天治愈（15天潜伏+10天治疗） 
# 隔离确诊患者：gamma = 1/15，按最长15天发病确诊后被隔离
# 隔离疑似人员：gamma = 1/3，平均3天被隔离
gamma = [1/25, 1/15, 1/3]
for i in range(3):
    #请使用SIR模型分别计算出未实施隔离，隔离确诊患者，隔离疑似人员的结果y，方法使用改进的欧拉法，并进行输出
    ########### Begin ###########
    simulation = SIR(beta=beta, gamma=gamma[i], y0=y0)
    y = simulation.solve(t, 'T')
    print(y)
    ########### End ###########   
    
    I = y[:, 1]
    plt.plot(t, I)

plt.legend(['未实施隔离', '确诊隔离', '疑似隔离'])
plt.title('隔离与为隔离机制下新型冠状病毒发展趋势对比分析')
plt.xlabel('时间（天）')
plt.ylabel('人数')
Vx = [1.0e6, 5.0e6] + [i * 1.e7 for i in range(11)]
plt.yticks(Vx, ['%d'%e for e in Vx])

plt.savefig('test2/test2.png')