SISSynchronousDynamicsRewireNeighbour.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from GraphWithSynchronousDynamics import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SISSynchronousDynamicsRewireNeighbour(GraphWithSynchronousDynamics):\n",
    "    '''A graph with a particular SIR dynamics. We use probabilities\n",
    "    to express infection and recovery per timestep, and run the system\n",
    "    using synchronous dynamics.'''\n",
    "    \n",
    "    # the possible dynamics states of a node for SIR dynamics\n",
    "    SUSCEPTIBLE = 'susceptible'\n",
    "    INFECTED = 'infected' \n",
    "    NEIGHBOURHOODS = dict()\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}\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",
    "        GraphWithSynchronousDynamics.__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",
    "        \n",
    "            \n",
    "    def _before( self ):\n",
    "        '''Seed the network with infected nodes, and mark all edges\n",
    "        as unoccupied by the dynamics.'''\n",
    "        self._infected = []       # in case we re-run from a dirty intermediate state\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",
    "        for (n, m, data) in self.edges_iter(data = True):\n",
    "            data[self.OCCUPIED] = False\n",
    "\n",
    "    def _after(self):\n",
    "        '''No processing to do after completion'''\n",
    "        pass\n",
    "            \n",
    "    def _dynamics_step( self ):\n",
    "        '''Optimised per-step dynamics that only runs the dynamics at infected\n",
    "        nodes, since they're the only places where state changes originate. At the\n",
    "        end of each timestep we re-build the infected node list.\n",
    "        returns: the number of events that happened in this timestep'''\n",
    "        events = 0\n",
    "        \n",
    "        # run model dynamics on all infected nodes\n",
    "        for n in self.POPULATION[self.INFECTED]:\n",
    "            events += self.model_infect_recover(n)\n",
    "        \n",
    "        # run rewire dynamics on all susceptible nodes\n",
    "        for s in self.POPULATION[self.SUSCEPTIBLE]:\n",
    "            events += self.model_rewire(s)\n",
    "        \n",
    "        return events\n",
    "            \n",
    "    def model_infect_recover( self, node_selected ):\n",
    "        '''Apply the SIR dynamics to node n. From the re-definition of dynamics_step()\n",
    "        we already know this node is infected.\n",
    "\n",
    "        n: the node\n",
    "        returns: the number of changes made'''\n",
    "        events = 0\n",
    "        \n",
    "        # infect susceptible neighbours with probability pInfect\n",
    "        for (_, neighbour, data) in self.edges_iter(node_selected, data = True):\n",
    "            if self.node[neighbour][self.DYNAMICAL_STATE] is self.SUSCEPTIBLE:\n",
    "                if numpy.random.random() <= self.p_infect:\n",
    "                    events += 1\n",
    "                    \n",
    "                    # infect the node\n",
    "                    self.update_node(neighbour,self.SUSCEPTIBLE,self.INFECTED)\n",
    "                        \n",
    "                    # label the edge we traversed as occupied\n",
    "                    data[self.OCCUPIED] = True\n",
    "    \n",
    "        # recover with probability pRecover\n",
    "        if numpy.random.random() <= self.p_recover:\n",
    "            # recover the node\n",
    "            events = events + 1\n",
    "            self.update_node(node_selected,self.INFECTED,self.SUSCEPTIBLE)\n",
    "                \n",
    "        return events\n",
    "    \n",
    "    def model_rewire( self, node_selected ):\n",
    "        '''Apply the rewire dynamics to node n. From the re-definition of dynamics_step()\n",
    "        we already know this node is susceptible.\n",
    "\n",
    "        n: the node\n",
    "        returns: the number of changes made'''\n",
    "        events = 0\n",
    "        \n",
    "        # rewire from infected neighbours with probability p_rewire\n",
    "        for (_, neighbour, data) in self.edges_iter(node_selected, data = True):\n",
    "            # If neighbour is infected\n",
    "            if self.node[neighbour][self.DYNAMICAL_STATE] is self.INFECTED:\n",
    "                # Check if rewire occurs\n",
    "                if numpy.random.random() <= self.p_rewire:\n",
    "                    # Try to rewire\n",
    "                    if (self.rewire(node_selected, neighbour)):\n",
    "                        events += 1\n",
    "                    \n",
    "        return events\n",
    "    \n",
    "    def at_equilibrium( self ):\n",
    "        '''SIR dynamics is at equilibrium if there are no more\n",
    "        infected nodes left in the network or if we've exceeded\n",
    "        the default simulation length.\n",
    "        \n",
    "        returns: True if the model has stopped'''\n",
    "        self.STATISTICS['peak_infection'] = max(self._pop_dist[self.INFECTED].iteritems(), key=operator.itemgetter(1))\n",
    "            \n",
    "        if self.CURRENT_TIMESTEP >= self._time_limit:\n",
    "            return True\n",
    "        else:\n",
    "            return (len(self.POPULATION[self.INFECTED]) == 0)\n",
    "        \n",
    "    def rewire(self, node = 0 , neighbour = 0):\n",
    "        '''Rewires the link between node and neighbour. Removes the link and tries to find a\n",
    "        valid node to connect to. Returns true if node found, false otherwise.'''\n",
    "        \n",
    "        self.remove_edges_from([(node, neighbour)])\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[node])) 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[node] = []\n",
    "            for j in self.neighbors(node):\n",
    "                self.NEIGHBOURHOODS[node].extend(self.neighbors(j))\n",
    "            self.NEIGHBOURHOODS[node] = [c for c in list(set(self.NEIGHBOURHOODS[node])) if c != node]\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 node not in self.NEIGHBOURHOODS[k]:\n",
    "                    self.NEIGHBOURHOODS[k].append(node)\n",
    "            \n",
    "            self.add_edge(node, candidates[new_neighb_index])\n",
    "            return True\n",
    "        \n",
    "        return False"
   ]
  }
 ],
 "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
}