SISStochasticDynamics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from GraphWithStochasticDynamics import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SISStochasticDynamics(GraphWithStochasticDynamics):\n",
    "    '''An SIS dynamics with stochastic simulation.'''\n",
    "\n",
    "    # the possible dynamics states of a node for SIR dynamics\n",
    "    SUSCEPTIBLE = 'susceptible'\n",
    "    INFECTED = 'infected'\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 ):\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",
    "        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",
    "        \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",
    "        if self.CURRENT_TIMESTEP >= self._time_limit:\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.SUSCEPTIBLE)\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",
    "        # add all the edges incident on this node connected to infected nodes\n",
    "        for (_, mp, datap) in self.edges_iter(n, data = True):\n",
    "            if self.node[mp][self.DYNAMICAL_STATE] == self.INFECTED:\n",
    "                self._si.insert(0, (mp, n, datap))\n",
    "                \n",
    "        \n",
    "    def transitions( self, t ):\n",
    "        '''Return the transition vector for the dynamics.\n",
    "        \n",
    "        t: time (ignored)\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()) ]"
   ]
  }
 ],
 "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
}