import numpy as np import matplotlib.pyplot as plt from scipy import integrate import sys N = 100000. I0 = 10. beta = 0.3 mu = 0.1 deltaT = 0.1 # Question 1 def explicit(N, I0, deltaT): S = [N-I0] I = [I0] R = [0] for i in range(int(365/deltaT)): S1 = S[-1]*(1-beta*I[-1]*deltaT/N) I1 = I[-1]*(1-mu*deltaT+beta*S[-1]*deltaT/N) R1 = R[-1]+mu*deltaT*I[-1] S.append(S1) I.append(I1) R.append(R1) return [S, I, R] # Question 2 S, I, R = explicit(N, I0, deltaT) # Question 3 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 4 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)) # Question 5 plt.plot(time, SIR) time = np.arange(0, 365+deltaT, deltaT) plt.plot(time, S, '+') plt.plot(time, I, '+') plt.plot(time, R, '+') plt.legend(('S', 'I', 'R', 'Se', 'Ie', 'Re')) plt.savefig('comparison.png') plt.close() # Question 6 def error(I, Ie, deltaT): sumI = np.sum(I) sumIe = np.sum(Ie)*deltaT return np.abs(sumIe-sumI)/365 # Question 7 deltas = np.linspace(0.001, 1, 10) err = [] for delta in deltas: print >> sys.stderr, delta S, I, R = explicit(N, I0, delta) err.append(error(SIR.T[1], I, delta)) # Question 8 plt.plot(deltas, err) plt.savefig('errorSIR.png') plt.close()