SIS model example - synchronous updating (mostly borrowed from PyCX)

Import packages, suppress matplotlib warnings

import pycxsimulator
from pylab import *
import networkx as nx
import warnings
warnings.filterwarnings('ignore')

Set up some basic parameters, etc.

p_i = 0.5 # infection probability per contact
p_r = 0.5 # recovery probability

prev = [] # we will store the disease prevalence over time here

Define our initialize function---for now we'll use the Karate Club graph, but you can use whatever network you like!

def initialize():
    global g, nextg, prev
    g = nx.karate_club_graph()
    # g = nx.barabasi_albert_graph(50,5)
    g.pos = nx.spring_layout(g)
    nx.set_node_attributes(g, 0, 'state')  # everyone starts off susceptible
    g.node[1]['state'] = 1                 # set one node to be infected (index case)
    nextg = g.copy()
    nextg.pos = g.pos
    prev.append(1/len(g.nodes)) # initial prevalence

Define our update function and track prevalence as we go:

def update():
    global g, nextg, prev
    curprev = 0
    for a in g.nodes:
        if g.node[a]['state'] == 0: # if susceptible
            nextg.node[a]['state'] = 0
            for b in g.neighbors(a):
                if g.node[b]['state'] == 1: # if neighbor b is infected
                    if random() < p_i:
                        nextg.node[a]['state'] = 1
        else: # if infected
            curprev += 1
            nextg.node[a]['state'] = 0 if random() < p_r else 1
    prev.append(curprev/len(g.nodes()))
    g = nextg

Define our observation:

def observe():
    global g, prev
    cla()
    nx.draw(g, cmap = cm.Wistia, vmin = 0, vmax = 1,
            node_color = [g.node[i]['state'] for i in g.nodes],
            pos = g.pos)

And off we go!

pycxsimulator.GUI().start(func=[initialize, observe, update])

Plot the epidemic curve:

epicurve = scatter(range(len(prev)), prev)
xlabel("Time")
ylabel("Prevalence")
show(epicurve)