Skip to content

Instantly share code, notes, and snippets.

@tigercosmos

tigercosmos/SIR.py

Created Mar 31, 2020
Embed
What would you like to do?
# -------------------------------------------------------------------------------
# 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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.