Skip to content

Instantly share code, notes, and snippets.

@tigercosmos

tigercosmos/agent-SIR.py

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
You can’t perform that action at this time.