SIR.py

# -------------------------------------------------------------------------------
# SIR Model
# Author:       Liu An Chi @tigercosmos
# -------------------------------------------------------------------------------
from pylab import *
from scipy import interpolate

# Parameters
Beta = 1
Gamma = 0.5
N = 1000
day = 20
I0 = 10


def SIR(S, I, R):

    dS = - Beta*S*I/N
    dR = Gamma*I
    dI = - dS - dR
    s = S + dS
    i = I + dI
    r = R + dR
    if (s < 0):
        s = 0
    if (i > N):
        i = N
    if (r > N):
        r = N

    return s, i, r


if __name__ == '__main__':
    I = I0
    S = N - I0
    R = 0

    LS = [S]
    LI = [I]
    LR = [R]

    for i in range(0, day):
        S, I, R = SIR(S, I, R)

        LS.append(S)
        LI.append(I)
        LR.append(R)

    # Draw the figure
    t = arange(0, day+1, 1)
    ls = interpolate.InterpolatedUnivariateSpline(t, LS)
    li = interpolate.InterpolatedUnivariateSpline(t, LI)
    lr = interpolate.InterpolatedUnivariateSpline(t, LR)
    tnew = arange(0.001, day, 1/20)
    lsnew = ls(tnew)
    linew = li(tnew)
    lrnew = lr(tnew)
    line1, = plot(tnew, lsnew, label='S')
    line2, = plot(tnew, linew, label='I')
    line3, = plot(tnew, lrnew, label='R')

    legend(handles=[line1, line2, line3],
           shadow=True, loc=(0.85, 0.4))  # handle

    line11, = plot(t, LS, '.')
    line22, = plot(t, LI, '.')
    line33, = plot(t, LR, '.')

    text(16.5, 240, 'Beta=%g' % (Beta))
    text(16.5, 190, 'Gamma=%g' % (Gamma))
    text(16.5, 150, 'Infected:%d' % (I0))
    v = [0, 20, 0, 1000]
    axis(v)
    xlabel('time (days)')
    ylabel('Population(people)')
    title('SIR Model')
    grid(True)
    show()