Skip to content

Instantly share code, notes, and snippets.

@tigercosmos
Created March 31, 2020 17:59
Show Gist options
  • Save tigercosmos/9f0251283bda230aafe8582305cc71f1 to your computer and use it in GitHub Desktop.
Save tigercosmos/9f0251283bda230aafe8582305cc71f1 to your computer and use it in GitHub Desktop.
# -------------------------------------------------------------------------------
# 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