agent-SIR.py

# -------------------------------------------------------------------------------
# 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()