"""
Simple version of the SIR model, formulated
and solved as a system of difference equations. 
The same system can be derived by applying the
Forward Euler method to the ODE-based version
of the model. 
"""


import numpy as np
import matplotlib.pyplot as plt

beta = 1.44
nu = 0.2                #1/nu = recovery time (5 days)
dt = 0.1                # time measured in days
D = 30                  # simulate for D days
steps = int(D/dt)       

t = np.linspace(0, steps*dt, steps+1)
S = np.zeros(steps+1)
I = np.zeros(steps+1)
R = np.zeros(steps+1)

S[0] = 50
I[0] = 1
N = S[0]+I[0]  #total population

for n in range(steps):
    S[n+1] = S[n] - dt*beta*S[n]*I[n]/N
    I[n+1] = I[n] + dt*beta*S[n]*I[n]/N - dt*nu*I[n]
    R[n+1] = R[n] + dt*nu*I[n]

# Plot the graphs
plt.plot(t, S, 'k-', t, I, 'b-', t, R, 'r-')
plt.legend(['S', 'I', 'R'], loc='lower right')
plt.xlabel('Time (days)')
plt.show()