import numpy as np import matplotlib.pyplot as plt from scipy import integrate N = 100000. I0 = 10. beta = 0.3 mu = 0.1 # Question 1 def new_cases_SIR(pop, t, beta, mu, N): S = pop[0] I = pop[1] R = pop[2] dS = -beta*S*I/N dI = beta*S*I/N-mu*I dR = mu*I return (dS, dI, dR) # Question 2 initial_pop = (N-I0, I0, 0) time = np.arange(0, 365, 1) SIR = integrate.odeint(new_cases_SIR, initial_pop, time, args = (beta, mu, N)) plt.plot(time, SIR) plt.legend(('S','I','R')) plt.savefig('SIR.png') plt.close() # Question 3 + 4 betas = np.arange(0.15, 0.31, 0.01) maxI = [] maxR = [] for beta in betas: SIR = integrate.odeint(new_cases_SIR, initial_pop, time, args = (beta, mu, N)) maxI.append(max(SIR.T[1])) maxR.append(max(SIR.T[2])) plt.plot(time, SIR.T[1]) plt.legend(betas) plt.savefig('I-R0.png') plt.close() # Question 5 plt.plot(betas/mu, maxI) plt.plot(betas/mu, maxR) plt.legend(('maxI', 'maxR')) plt.savefig('max.png') plt.close() # Question 6 def new_cases_SIS(pop, t, beta, mu, N): S = pop[0] I = pop[1] dS = -beta*S*I/N+mu*I dI = beta*S*I/N-mu*I return (dS, dI) initial_pop = (N-I0, I0) beta = 0.3 SIS = integrate.odeint(new_cases_SIS, initial_pop, time, args = (beta, mu, N)) plt.plot(time, SIS) plt.legend(('S','I')) plt.savefig('SIS.png') plt.close()