Skip to content

Instantly share code, notes, and snippets.

@tigercosmos
Last active Mar 31, 2020
Embed
What would you like to do?
# -------------------------------------------------------------------------------
# Agent-based SIR Model
# Author: Liu An Chi @tigercosmos
# -------------------------------------------------------------------------------
from pylab import *
from scipy import interpolate
from random import random
from math import log
from enum import Enum
# Exponential Distribution
def InfectiousEvent(rate, dt):
try:
return -log(random()) / rate*1000 < dt
except ZeroDivisionError:
return 0
def RecoveredEvent(rate, dt):
try:
return -log(random()) / rate*1.25 < dt
except ZeroDivisionError:
return 0
# Parameters
Beta = 1
Gamma = 0.5
N = 1000
day = 20
I0 = 10
class State(Enum):
SUS = 1
INF = 2
RECV = 3
class Person:
def __init__(self, id, state):
self.id = id
self.state = state
def step(self, people, dt=1):
if self.state == State.SUS:
new_beta = 0
# suppose that touch all people,
# and add all's beta as new beta
for person in people:
if person.state == State.INF:
new_beta += Beta
if InfectiousEvent(new_beta, dt):
self.state = State.INF
elif self.state == State.INF:
if RecoveredEvent(Gamma, dt):
self.state = State.RECV
def sir_count(people):
s = 0
i = 0
r = 0
for p in people:
if p.state == State.SUS:
s += 1
elif p.state == State.INF:
i += 1
else:
r += 1
return s, i, r
if __name__ == '__main__':
LS = []
LI = []
LR = []
people = []
for i in range(0, N - I0):
people.append(Person(i, State.SUS))
for i in range(0, I0):
people.append(Person(i, State.INF))
for day in range(0, day + 1):
s, i, r = sir_count(people)
LS.append(s)
LI.append(i)
LR.append(r)
print(s, i, r)
for person in people:
person.step(people)
# 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, day, 0, N]
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