#! /usr/bin/env python
# -*- coding: utf-8 -*-
# This simple example on how to do animations using graph-tool. Here we do a
# simple simulation of an S->I->R->S epidemic model, where each vertex can be in
# one of the following states: Susceptible (S), infected (I), recovered (R). A
# vertex in the S state becomes infected either spontaneously with a probability
# 'x' or because a neighbour is infected. An infected node becomes recovered
# with probability 'r', and a recovered vertex becomes again susceptible with
# probability 's'.
# DISCLAIMER: The following code is definitely not the most efficient approach
# if you want to simulate this dynamics for very large networks, and/or for very
# long times. The main purpose is simply to highlight the animation capabilities
# of graph-tool.
from graph_tool.all import *
from numpy.random import *
import sys, os, os.path
seed(42)
seed_rng(42)
# We need some Gtk and gobject functions
from gi.repository import Gtk, Gdk, GdkPixbuf
import gi._gobject as gobject
# We will use the network of network scientists, and filter out the largest
# component
g = collection.data["netscience"]
g = GraphView(g, vfilt=label_largest_component(g), directed=False)
g = Graph(g, prune=True)
pos = g.vp["pos"] # layout positions
# We will filter out vertices which are in the "Recovered" state, by masking
# them using a property map.
removed = g.new_vertex_property("bool")
# SIRS dynamics parameters:
x = 0.001 # spontaneous outbreak probability
r = 0.1 # I->R probability
s = 0.01 # R->S probability
# (Note that the S->I transition happens simultaneously for every vertex with a
# probability equal to the fraction of non-recovered neighbours which are
# infected.)
# The states would usually be represented with simple integers, but here we will
# use directly the color of the vertices in (R,G,B,A) format.
S = [1, 1, 1, 1] # White color
I = [0, 0, 0, 1] # Black color
R = [0.5, 0.5, 0.5, 1.] # Grey color (will not actually be drawn)
# Initialize all vertices to the S state
state = g.new_vertex_property("vector")
for v in g.vertices():
state[v] = S
# Newly infected nodes will be highlighted in red
newly_infected = g.new_vertex_property("bool")
# If True, the frames will be dumped to disk as images.
offscreen = sys.argv[1] == "offscreen" if len(sys.argv) > 1 else False
max_count = 500
if offscreen and not os.path.exists("./frames"):
os.mkdir("./frames")
# This creates a GTK+ window with the initial graph layout
if not offscreen:
win = GraphWindow(g, pos, geometry=(500, 400),
edge_color=[0.6, 0.6, 0.6, 1],
vertex_fill_color=state,
vertex_halo=newly_infected,
vertex_halo_color=[0.8, 0, 0, 0.6])
else:
count = 0
win = Gtk.OffscreenWindow()
win.set_default_size(500, 400)
win.graph = GraphWidget(g, pos,
edge_color=[0.6, 0.6, 0.6, 1],
vertex_fill_color=state,
vertex_halo=newly_infected,
vertex_halo_color=[0.8, 0, 0, 0.6])
win.add(win.graph)
# This function will be called repeatedly by the GTK+ main loop, and we use it
# to update the state according to the SIRS dynamics.
def update_state():
newly_infected.a = False
removed.a = False
# visit the nodes in random order
vs = list(g.vertices())
shuffle(vs)
for v in vs:
if state[v] == I:
if random() < r:
state[v] = R
elif state[v] == S:
if random() < x:
state[v] = I
else:
ns = list(v.out_neighbours())
if len(ns) > 0:
w = ns[randint(0, len(ns))] # choose a random neighbour
if state[w] == I:
state[v] = I
newly_infected[v] = True
elif random() < s:
state[v] = S
if state[v] == R:
removed[v] = True
# Filter out the recovered vertices
g.set_vertex_filter(removed, inverted=True)
# The following will force the re-drawing of the graph, and issue a
# re-drawing of the GTK window.
win.graph.regenerate_surface(lazy=False)
win.graph.queue_draw()
# if doing an offscreen animation, dump frame to disk
if offscreen:
global count
pixbuf = win.get_pixbuf()
pixbuf.savev(r'./frames/sirs%06d.png' % count, 'png', [], [])
if count > max_count:
sys.exit(0)
count += 1
# We need to return True so that the main loop will call this function more
# than once.
return True
# Bind the function above as an 'idle' callback.
cid = gobject.idle_add(update_state)
# We will give the user the ability to stop the program by closing the window.
win.connect("delete_event", Gtk.main_quit)
# Actually show the window, and start the main loop.
win.show_all()
Gtk.main()