SIRAsynchronousDynamics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from GraphWithAsynchronousDynamics import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SIRAsynchronousDynamics(GraphWithAsynchronousDynamics):\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",
    "    RECOVERED = 'recovered'\n",
    "    \n",
    "    # list of infected nodes, the sites of all the dynamics\n",
    "    _infected = []\n",
    "    _susceptible = []\n",
    "    _recovered = []\n",
    "    \n",
    "    \n",
    "    def __init__( self, time_limit = 10000, p_infect = 0.0, p_recover = 1.0, p_infected = 0.0, graph = None ):\n",
    "        '''Generate a graph with dynamics for the given parameters.\n",
    "        \n",
    "        pInfect: infection probability (defaults to 0.0)\n",
    "        pRecover: probability of recovery (defaults to 1.0)\n",
    "        pInfected: initial infection probability (defaults to 0.0)\n",
    "        g: 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",
    "        GraphWithAsynchronousDynamics.__init__(self, time_limit = time_limit, graph = graph, states = states, rates = rates)\n",
    "        self._p_infect = p_infect\n",
    "        self._p_recover = p_recover\n",
    "        self._p_infected = p_infected\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",
    "        pass\n",
    "    \n",
    "    def set_timestep_rate(self):\n",
    "        '''Calculate the minimum time-jump. Choose the maximum from the recovery rate\n",
    "        or the max node degree * the infection rate. Timestep will be 1 over this value'''\n",
    "        max_degree = max(self.degree().values())\n",
    "        return max(max_degree*self._p_infect, self._p_recover)\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",
    "        \n",
    "        '''if(t%1000 == 0.0):\n",
    "            print(\"Processing timestep \",t)'''\n",
    "        \n",
    "        if self.CURRENT_TIMESTEP >= self._time_limit:\n",
    "            return True\n",
    "        else:\n",
    "            return (len(self.POPULATION[self.INFECTED]) == 0)\n",
    "            \n",
    "    \n",
    "    def asyn_node_action(self, node = 0, dt = 0.0, r = 0.0):\n",
    "        \n",
    "        # If chosen node is infected\n",
    "        if(self.node[node][self.DYNAMICAL_STATE] == self.INFECTED):\n",
    "            # If random number lower than Number of nodes x rate recovery * timestep\n",
    "            if(r < self.order()*self._p_recover*dt):\n",
    "                # Node recovers\n",
    "                self.update_node(node,self.INFECTED,self.RECOVERED)\n",
    "                return True\n",
    "        # If chosen node is susceptible\n",
    "        elif(self.node[node][self.DYNAMICAL_STATE] == self.SUSCEPTIBLE):\n",
    "            infected_neighbours = 0\n",
    "            # Count neighbours - TODO - better way to do this? Have to count the neighbours each time\n",
    "            for (_, mp, datap) in self.edges_iter(node, data = True):\n",
    "                if self.node[mp][self.DYNAMICAL_STATE] == self.INFECTED:\n",
    "                        infected_neighbours += 1\n",
    "            # If random number lower than Number of nodes x rate infection * timestep * infected_neighbours\n",
    "            if(r < self.order()*self._p_infect*dt*infected_neighbours):\n",
    "                # Node gets infected\n",
    "                self.update_node(node,self.SUSCEPTIBLE,self.INFECTED)\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
}