SRI Immunization Model

The Code

Here's the code for the model relying only on numpy, networkx, holoviews and matplotlib in the background.

In [1]:
import collections
import itertools
import math

import numpy as np
np.seterr(divide='ignore')
import numpy.random as rnd
import networkx as nx

import param
import holoviews as hv

SPREADING_SUSCEPTIBLE = 'S'
SPREADING_VACCINATED = 'V'
SPREADING_INFECTED = 'I'
SPREADING_RECOVERED = 'R'
DEAD = 'D'

class SRI_Model(param.Parameterized):
    """
    Implementation of the SRI epidemiology model
    using NetworkX and HoloViews for visualization.
    This code has been adapted from Simon Dobson's
    code here:
    
    http://www.simondobson.org/complex-networks-complex-processes/epidemic-spreading.html
    
    In addition to his basic parameters I've added
    additional states to the model, a node may be
    in one of the following states:
    
      * Susceptible: Can catch the disease from a connected node.
      * Vaccinated: Immune to infection.
      * Infected: Has the disease and may pass it on to any connected node.
      * Recovered: Immune to infection.
      * Dead: Edges are removed from graph.
    """

    network = param.ClassSelector(class_=nx.Graph, default=None, doc="""
        A custom NetworkX graph, instead of the default Erdos-Renyi graph.""")
    
    visualize = param.Boolean(default=True, doc="""
        Whether to compute layout of network for visualization.""")
    
    # Initial parameters
    
    N = param.Integer(default=1000, doc="""
        Number of nodes to simulate.""")
    
    mean_connections = param.Number(default=10, doc="""
        Mean number of connections to make to other nodes.""")
    
    pSick = param.Number(default=0.01, doc="""
        Probability of a node to be initialized in sick state.""", bounds=(0, 1))

    pVaccinated = param.Number(default=0.1, bounds=(0, 1), doc="""
        Probability of a node to be initialized in vaccinated state.""")
    
    # Simulation parameters
    
    pInfect = param.Number(default=0.3, doc="""
        Probability of infection on each time step.""", bounds=(0, 1))
    
    pRecover = param.Number(default=0.05, doc="""
        Probability of recovering if infected on each timestep.""", bounds=(0, 1))
    
    pDeath = param.Number(default=0.1, doc="""
        Probability of death if infected on each timestep.""", bounds=(0, 1))
    
    def __init__(self, **params):
        super(SRI_Model, self).__init__(**params)
        if not self.network:
            self.g = nx.erdos_renyi_graph(self.N, float(self.mean_connections)/self.N)
        else:
            self.g = self.network
        self.vaccinated, self.infected = self.spreading_init()
        self.model = self.spreading_make_sir_model()
        self.color_mapping = [SPREADING_SUSCEPTIBLE,
                              SPREADING_VACCINATED,
                              SPREADING_INFECTED,
                              SPREADING_RECOVERED, DEAD]
        if self.visualize:
            self.pos = nx.spring_layout(self.g, iterations = 50,
                                        k = 2/(math.sqrt(self.g.order())))

    def spreading_init(self):
        """Initialise the network with vaccinated, susceptible and infected states."""
        vaccinated, infected = 0, []
        for i in self.g.node.keys():
            self.g.node[i]['transmissions'] = 0
            if(rnd.random() <= self.pVaccinated): 
                self.g.node[i]['state'] = SPREADING_VACCINATED
                vaccinated += 1
            elif(rnd.random() <= self.pSick):
                self.g.node[i]['state'] = SPREADING_INFECTED
                infected.append(i)
            else:
                self.g.node[i]['state'] = SPREADING_SUSCEPTIBLE
        return vaccinated, infected

    def spreading_make_sir_model(self):
        """Return an SIR model function for given infection and recovery probabilities."""
        # model (local rule) function
        def model( g, i ):
            if g.node[i]['state'] == SPREADING_INFECTED:
                # infect susceptible neighbours with probability pInfect
                for m in g.neighbors(i):
                    if g.node[m]['state'] == SPREADING_SUSCEPTIBLE:
                        if rnd.random() <= self.pInfect:
                            g.node[m]['state'] = SPREADING_INFECTED
                            self.infected.append(m)
                            g.node[i]['transmissions'] += 1

                # recover with probability pRecover
                if rnd.random() <= self.pRecover:
                    g.node[i]['state'] = SPREADING_RECOVERED
                elif rnd.random() <= self.pDeath:
                    edges = [edge for edge in self.g.edges() if i in edge] 
                    g.node[i]['state'] = DEAD
                    g.remove_edges_from(edges)

        return model

    def step(self):
        """Run a single step of the model over the graph."""
        for i in self.g.node.keys():
            self.model(self.g, i)

    def run(self, steps):
        """
        Run the network for the specified number of time steps
        """
        for i in range(steps):
            self.step()

    def network_data(self):
        """
        Return the network edge paths and node positions,
        requires visualize parameter to be enabled.
        """
        if not self.visualize:
            raise Exception("Enable visualize option to get network data.")

        nodeMarkers = []
        overlay = []
        points = np.array([self.pos[v] for v in self.g.nodes_iter()])
        paths = []
        for e in self.g.edges_iter():
            xs = [ self.pos[e[0]][0], self.pos[e[1]][0] ]
            ys = [ self.pos[e[0]][1], self.pos[e[1]][1] ]
            paths.append(np.array(zip(xs, ys)))
        return paths, points

    def stats(self):
        """
        Return an ItemTable with statistics on the network data.
        """
        state_labels = hv.OrderedDict([('S', 'Susceptible'), ('V', 'Vaccinated'), ('I', 'Infected'),
                                    ('R', 'Recovered'), ('D', 'Dead')])
        counts = collections.Counter()
        transmissions = []
        for n in self.g.nodes_iter():
            state = state_labels[self.g.node[n]['state']]
            counts[state] += 1
            if n in self.infected:
                transmissions.append(self.g.node[n]['transmissions'])
        data = hv.OrderedDict([(l, counts[l])
                               for l in state_labels.values()])
        
        infected = len(set(self.infected))
        unvaccinated = float(self.N-self.vaccinated)
        
        data['$R_0$'] = np.mean(transmissions) if transmissions else 0
        data['Death rate DR'] = np.divide(float(data['Dead']),self.N)
        data['Infection rate IR'] = np.divide(float(infected), self.N)
        if unvaccinated:
            unvaccinated_dr = data['Dead']/unvaccinated
            unvaccinated_ir = infected/unvaccinated
        else:
            unvaccinated_dr = 0
            unvaccinated_ir = 0
        data['Unvaccinated DR'] = unvaccinated_dr
        data['Unvaccinated IR'] = unvaccinated_ir
        return hv.ItemTable(data)

    def animate(self, steps):
        """
        Run the network for the specified number of steps accumulating animations
        of the network nodes and edges changing states and curves tracking the
        spread of the disease.
        """
        if not self.visualize:
            raise Exception("Enable visualize option to get compute network visulizations.")

        # Declare HoloMap for network animation and counts array
        network_hmap = hv.HoloMap(key_dimensions=['Time'])
        sird = np.zeros((steps, 5))
        
        # Declare dimensions and labels
        spatial_dims = [hv.Dimension('x', range=(-.1, 1.1)),
                        hv.Dimension('y', range=(-.1, 1.1))]
        state_labels = ['Susceptible', 'Vaccinated', 'Infected', 'Recovered', 'Dead']

        # Text annotation
        nlabel = hv.Text(0.9, 0.05, 'N=%d' % self.N)

        for i in range(steps):
            # Get path, point, states and count data
            paths, points = self.network_data()
            states = [self.color_mapping.index(self.g.node[n]['state'])
                      for n in self.g.nodes_iter()]
            state_array = np.array(states, ndmin=2).T
            (sird[i, :], _) = np.histogram(state_array, bins=list(range(6)))
            
            # Create network path and node Elements
            network_paths = hv.Path(paths, key_dimensions=spatial_dims)
            network_nodes = hv.Points(np.hstack([points, state_array]),
                                      key_dimensions=spatial_dims,
                                      value_dimensions=['State'])
            
            # Create overlay and accumulate in network HoloMap
            network_hmap[i] = (network_paths * network_nodes * nlabel).relabel(group='Network', label='SRI')
            self.step()

        # Create Overlay of Curves
        #extents = (-1, -1, steps, np.max(sird)+2)
        curves = hv.NdOverlay({label: hv.Curve(zip(range(steps), sird[:, i]),
                                               key_dimensions=['Time'], value_dimensions=['Count'])
                              for i, label in enumerate(state_labels)},
                              key_dimensions=[hv.Dimension('State', values=state_labels)])
        
        # Animate VLine on top of Curves
        distribution = hv.HoloMap({i: (curves * hv.VLine(i)).relabel(group='Counts', label='SRI')
                                   for i in range(steps)}, key_dimensions=['Time'])
        
        return network_hmap + distribution

The style

HoloViews allows use to define various style options in advance on the Store.options object.

In [2]:
hv.extension('bokeh', 'matplotlib')

# Set colors and style options for the Element types
from holoviews import Store, Options
from holoviews.core.options import Palette
opts = Store.options()

opts.Path      = Options('style', line_width=0.2, color='black')
opts.Curve     = Options('style', color=Palette('hot_r'))
opts.Histogram = Options('plot', show_grid=False)
opts.Overlay   = Options('plot', show_frame=False)
opts.HeatMap   = Options('plot', xrotation=90)
opts.ItemTable = Options('plot', width=900, height=50)

opts.Overlay.Network = Options('plot', xaxis=None, yaxis=None)
opts.Overlay.Counts  = Options('plot', show_grid=True)

opts.Points    = {'style': Options(cmap='hot_r', size=8, line_color='black'),
                  'plot':  Options(color_index=2)}
opts.VLine     = {'style': Options(color='black', line_width=1),
                  'plot':  Options(show_grid=True)}