SIRStochasticDynamicsRewireNeighbour.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from GraphWithStochasticDynamics import *\n",
    "import operator\n",
    "import csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class SIRStochasticDynamicsRewireNeighbour(GraphWithStochasticDynamics):\n",
    "    '''An SIR dynamics with stochastic simulation.'''\n",
    "\n",
    "    # the possible dynamics states of a node for SIR dynamics\n",
    "    SUSCEPTIBLE = 'susceptible'\n",
    "    INFECTED = 'infected'\n",
    "    RECOVERED = 'recovered'\n",
    "    NEIGHBOURHOODS = dict()\n",
    "    \n",
    "    # list of SI edges connecting a susceptible to an infected node\n",
    "    _si = []\n",
    "        \n",
    "    def __init__( self, time_limit = 10000, p_infect = 0.0, p_recover = 1.0, p_infected = 0.0, graph = None, p_rewire = 0.0 ):\n",
    "        '''Generate a graph with dynamics for the given parameters.\n",
    "        \n",
    "        p_infect: infection probability (defaults to 0.0)\n",
    "        p_recover: probability of recovery (defaults to 1.0)\n",
    "        p_infected: initial infection probability (defaults to 0.0)\n",
    "        graph: the graph to copy from (optional)'''\n",
    "        states = {self.SUSCEPTIBLE,self.INFECTED,self.RECOVERED}\n",
    "        rates = dict()\n",
    "        rates['p_infect'] = p_infect\n",
    "        rates['p_recover'] = p_recover\n",
    "        rates['p_infected'] = p_infected\n",
    "        rates['p_rewire'] = p_rewire\n",
    "        GraphWithStochasticDynamics.__init__(self, time_limit = time_limit, graph = graph, states = states, rates = rates)\n",
    "        self.p_infected = p_infected\n",
    "        self.p_infect = p_infect\n",
    "        self.p_recover = p_recover\n",
    "        self.p_rewire = p_rewire\n",
    "        \n",
    "        # Create a dictionary of all nodes 2 steps away from each node (i.e. neighbours of neighbours)\n",
    "        for i in self.nodes():\n",
    "            self.NEIGHBOURHOODS[i] = []\n",
    "            for j in self.neighbors(i):\n",
    "                self.NEIGHBOURHOODS[i].extend(self.neighbors(j))\n",
    "            self.NEIGHBOURHOODS[i] = [c for c in list(set(self.NEIGHBOURHOODS[i])) if c != i]\n",
    "        \n",
    "    def before( self ):\n",
    "        '''Seed the network with infected nodes, extract the initial set of\n",
    "        SI nodes, and mark all edges as unoccupied by the dynamics.'''\n",
    "        self._si = []\n",
    "        \n",
    "        # infect nodes\n",
    "        for n in self.node.keys():\n",
    "            if numpy.random.random() <= self.p_infected:\n",
    "                self.node[n][self.DYNAMICAL_STATE] = self.INFECTED\n",
    "            else:\n",
    "                self.node[n][self.DYNAMICAL_STATE] = self.SUSCEPTIBLE\n",
    "                \n",
    "        self.POPULATION = self.calculate_populations()        \n",
    "        \n",
    "        # extract the initial set of SI edges\n",
    "        for (n, m, data) in self.edges_iter(self.POPULATION[self.INFECTED], data = True):\n",
    "            if self.node[m][self.DYNAMICAL_STATE] == self.SUSCEPTIBLE:\n",
    "                self._si.insert(0, (n, m, data))\n",
    "        \n",
    "        # mark all edges as unoccupied\n",
    "        for (n, m, data) in self.edges_iter(data = True):\n",
    "            data[self.OCCUPIED] = False\n",
    "            \n",
    "    def after(self):\n",
    "        pass\n",
    "    \n",
    "    def at_equilibrium( self):\n",
    "        '''SIR dynamics is at equilibrium if there are no more infected nodes left\n",
    "        in the network, no susceptible nodes adjacent to infected nodes, or if we've\n",
    "        exceeded the default simulation length.\n",
    "        \n",
    "        returns: True if the model has stopped\n",
    "        \n",
    "        Extension: stop once the infection population drops below 90% of the current maximum\n",
    "        '''\n",
    "        if len(self.POPULATION[self.INFECTED]) < (0.9 * max(self._pop_dist[self.INFECTED].values())):\n",
    "            self.STATISTICS['peak_infection'] = max(self._pop_dist[self.INFECTED].iteritems(), key=operator.itemgetter(1))\n",
    "            return True\n",
    "        \n",
    "        if self.CURRENT_TIMESTEP >= self._time_limit:\n",
    "            self.STATISTICS['peak_infection'] = max(self._pop_dist[self.INFECTED].iteritems(), key=operator.itemgetter(1))\n",
    "            return True\n",
    "        else:\n",
    "            return ((len(self.POPULATION[self.INFECTED]) == 0))\n",
    "                    # or (len(self._si) == 0))\n",
    "\n",
    "    def infect( self ):\n",
    "        '''Infect a node chosen at random from the SI edges.'''\n",
    "         \n",
    "        # choose an SI edge\n",
    "        i = numpy.random.randint(len(self._si))\n",
    "        (n, m, data) = self._si[i]\n",
    "        \n",
    "        # infect the susceptible end\n",
    "        self.update_node(m,self.SUSCEPTIBLE,self.INFECTED)\n",
    "        \n",
    "        # label the edge we traversed as occupied\n",
    "        data[self.OCCUPIED] = True\n",
    "        \n",
    "        # remove all edges in the SI list from an infected node to this one\n",
    "        self._si = [ (np, mp, data) for (np, mp, data) in self._si if m != mp ]\n",
    "        \n",
    "        # add all the edges incident on this node connected to susceptible nodes\n",
    "        for (_, mp, datap) in self.edges_iter(m, data = True):\n",
    "            if self.node[mp][self.DYNAMICAL_STATE] == self.SUSCEPTIBLE:\n",
    "                self._si.insert(0, (m, mp, datap))\n",
    "        \n",
    "    def recover( self ):\n",
    "        '''Cause a node to recover.'''\n",
    "        \n",
    "        # choose an infected node at random\n",
    "        i = numpy.random.randint(len(self.POPULATION[self.INFECTED]))\n",
    "        n = self.POPULATION[self.INFECTED][i]\n",
    "        \n",
    "        # mark the node as recovered\n",
    "        self.update_node(n,self.INFECTED,self.RECOVERED)\n",
    "        \n",
    "        # remove all edges in the SI list incident on this node\n",
    "        self._si = [ (np, m, e) for (np, m, e) in self._si if np != n ]\n",
    "    \n",
    "    def rewire( self ):\n",
    "        '''Cause a node to rewire.'''\n",
    "       \n",
    "        # choose an SI edge\n",
    "        i = numpy.random.randint(len(self._si))\n",
    "        # remove the edge from the si list and from the overall graph structure\n",
    "        (n, m, data) = self._si.pop(i)\n",
    "        self.remove_edges_from([(n, m)])\n",
    "        \n",
    "        # Add a link to a new node\n",
    "        \n",
    "        # Candidate nodes are those 2 steps away (neighbours of neighbours) who are not infected\n",
    "        candidates = [c for c in list(set(self.NEIGHBOURHOODS[m])) if self.node[c][self.DYNAMICAL_STATE] != self.INFECTED]\n",
    "        \n",
    "        # Pick a candidate\n",
    "        if(len(candidates) >= 1):\n",
    "            new_neighb_index = numpy.random.randint(len(candidates))\n",
    "            \n",
    "            # Recalculate neighbourhoods\n",
    "            # Recalculate for susceptible node\n",
    "            self.NEIGHBOURHOODS[m] = []\n",
    "            for j in self.neighbors(m):\n",
    "                self.NEIGHBOURHOODS[m].extend(self.neighbors(j))\n",
    "            self.NEIGHBOURHOODS[m] = [c for c in list(set(self.NEIGHBOURHOODS[m])) if c != m]\n",
    "            \n",
    "            # Add suscpetible node to the neighbourhoods of those adjacent to the new node\n",
    "            for k in self.neighbors(candidates[new_neighb_index]):\n",
    "                if m not in self.NEIGHBOURHOODS[k]:\n",
    "                    self.NEIGHBOURHOODS[k].append(m)\n",
    "            \n",
    "            self.add_edge(m, candidates[new_neighb_index])\n",
    "            \n",
    "        \n",
    "    def transitions( self ):\n",
    "        '''Return the transition vector for the dynamics.\n",
    "        \n",
    "        returns: the transition vector'''\n",
    "        \n",
    "        # transitions are expressed as rates, whereas we're specified\n",
    "        # in terms of probabilities, so we convert the latter to the former.\n",
    "        return [ (len(self._si) * self.p_infect,        lambda t: self.infect()),\n",
    "                 (len(self.POPULATION[self.INFECTED]) * self.p_recover, lambda t: self.recover()),\n",
    "                 (len(self._si) * self.p_rewire,        lambda t: self.rewire())]\n",
    "            "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}