Skip to content

Commit

Permalink
Fixes to the SRI model notebook (#1699)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jlstevens @philippjfr
jlstevens authored and philippjfr committed Sep 29, 2017
1 parent 3680dcb commit 59ea7a2
Showing 1 changed file with 51 additions and 29 deletions.
80 changes: 51 additions & 29 deletions examples/topics/simulation/sri_model.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"import collections\n",
Expand All @@ -43,11 +45,6 @@
"import param\n",
"import holoviews as hv\n",
"\n",
"SPREADING_SUSCEPTIBLE = 'S'\n",
"SPREADING_VACCINATED = 'V'\n",
"SPREADING_INFECTED = 'I'\n",
"SPREADING_RECOVERED = 'R'\n",
"DEAD = 'D'\n",
"\n",
"class SRI_Model(param.Parameterized):\n",
" \"\"\"\n",
Expand Down Expand Up @@ -75,8 +72,6 @@
" visualize = param.Boolean(default=True, doc=\"\"\"\n",
" Whether to compute layout of network for visualization.\"\"\")\n",
" \n",
" # Initial parameters\n",
" \n",
" N = param.Integer(default=1000, doc=\"\"\"\n",
" Number of nodes to simulate.\"\"\")\n",
" \n",
Expand All @@ -89,8 +84,6 @@
" pVaccinated = param.Number(default=0.1, bounds=(0, 1), doc=\"\"\"\n",
" Probability of a node to be initialized in vaccinated state.\"\"\")\n",
" \n",
" # Simulation parameters\n",
" \n",
" pInfect = param.Number(default=0.3, doc=\"\"\"\n",
" Probability of infection on each time step.\"\"\", bounds=(0, 1))\n",
" \n",
Expand All @@ -100,6 +93,13 @@
" pDeath = param.Number(default=0.1, doc=\"\"\"\n",
" Probability of death if infected on each timestep.\"\"\", bounds=(0, 1))\n",
" \n",
" \n",
" SPREADING_SUSCEPTIBLE = 'S'\n",
" SPREADING_VACCINATED = 'V'\n",
" SPREADING_INFECTED = 'I'\n",
" SPREADING_RECOVERED = 'R'\n",
" DEAD = 'D'\n",
"\n",
" def __init__(self, **params):\n",
" super(SRI_Model, self).__init__(**params)\n",
" if not self.network:\n",
Expand All @@ -108,10 +108,10 @@
" self.g = self.network\n",
" self.vaccinated, self.infected = self.spreading_init()\n",
" self.model = self.spreading_make_sir_model()\n",
" self.color_mapping = [SPREADING_SUSCEPTIBLE,\n",
" SPREADING_VACCINATED,\n",
" SPREADING_INFECTED,\n",
" SPREADING_RECOVERED, DEAD]\n",
" self.color_mapping = [self.SPREADING_SUSCEPTIBLE,\n",
" self.SPREADING_VACCINATED,\n",
" self.SPREADING_INFECTED,\n",
" self.SPREADING_RECOVERED, self.DEAD]\n",
" if self.visualize:\n",
" self.pos = nx.spring_layout(self.g, iterations = 50,\n",
" k = 2/(math.sqrt(self.g.order())))\n",
Expand All @@ -122,34 +122,34 @@
" for i in self.g.node.keys():\n",
" self.g.node[i]['transmissions'] = 0\n",
" if(rnd.random() <= self.pVaccinated): \n",
" self.g.node[i]['state'] = SPREADING_VACCINATED\n",
" self.g.node[i]['state'] = self.SPREADING_VACCINATED\n",
" vaccinated += 1\n",
" elif(rnd.random() <= self.pSick):\n",
" self.g.node[i]['state'] = SPREADING_INFECTED\n",
" self.g.node[i]['state'] = self.SPREADING_INFECTED\n",
" infected.append(i)\n",
" else:\n",
" self.g.node[i]['state'] = SPREADING_SUSCEPTIBLE\n",
" self.g.node[i]['state'] = self.SPREADING_SUSCEPTIBLE\n",
" return vaccinated, infected\n",
"\n",
" def spreading_make_sir_model(self):\n",
" \"\"\"Return an SIR model function for given infection and recovery probabilities.\"\"\"\n",
" # model (local rule) function\n",
" def model( g, i ):\n",
" if g.node[i]['state'] == SPREADING_INFECTED:\n",
" if g.node[i]['state'] == self.SPREADING_INFECTED:\n",
" # infect susceptible neighbours with probability pInfect\n",
" for m in g.neighbors(i):\n",
" if g.node[m]['state'] == SPREADING_SUSCEPTIBLE:\n",
" if g.node[m]['state'] == self.SPREADING_SUSCEPTIBLE:\n",
" if rnd.random() <= self.pInfect:\n",
" g.node[m]['state'] = SPREADING_INFECTED\n",
" g.node[m]['state'] = self.SPREADING_INFECTED\n",
" self.infected.append(m)\n",
" g.node[i]['transmissions'] += 1\n",
"\n",
" # recover with probability pRecover\n",
" if rnd.random() <= self.pRecover:\n",
" g.node[i]['state'] = SPREADING_RECOVERED\n",
" g.node[i]['state'] = self.SPREADING_RECOVERED\n",
" elif rnd.random() <= self.pDeath:\n",
" edges = [edge for edge in self.g.edges() if i in edge] \n",
" g.node[i]['state'] = DEAD\n",
" g.node[i]['state'] = self.DEAD\n",
" g.remove_edges_from(edges)\n",
"\n",
" return model\n",
Expand Down Expand Up @@ -245,8 +245,9 @@
" state_array = np.array(states, ndmin=2).T\n",
" (sird[i, :], _) = np.histogram(state_array, bins=list(range(6)))\n",
" \n",
" path_data = [list(el[()]) for el in paths]\n",
" # Create network path and node Elements\n",
" network_paths = hv.Path(paths, key_dimensions=spatial_dims)\n",
" network_paths = hv.Path(path_data, key_dimensions=spatial_dims)\n",
" network_nodes = hv.Points(np.hstack([points, state_array]),\n",
" key_dimensions=spatial_dims,\n",
" value_dimensions=['State'])\n",
Expand All @@ -258,7 +259,7 @@
" # Create Overlay of Curves\n",
" #extents = (-1, -1, steps, np.max(sird)+2)\n",
" curves = hv.NdOverlay({label: hv.Curve(zip(range(steps), sird[:, i]),\n",
" key_dimensions=['Time'], value_dimensions=['Count'])\n",
" kdims=['Time'], vdims=['Count'])\n",
" for i, label in enumerate(state_labels)},\n",
" key_dimensions=[hv.Dimension('State', values=state_labels)])\n",
" \n",
Expand Down Expand Up @@ -338,7 +339,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"experiment1_params = dict(pInfect=0.08, pRecover=0.08, pSick=0.15,\n",
Expand Down Expand Up @@ -508,7 +511,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"experiment2_params = dict(N=1000, pInfect=0.05, pRecover=0.05,\n",
Expand All @@ -525,7 +530,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"exp2_dims = ['Connections', 'pVaccinated']\n",
Expand Down Expand Up @@ -574,7 +581,9 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"quantities = ['Unvaccinated IR', 'Infection rate IR', 'Death rate DR', '$R_0$']\n",
Expand Down Expand Up @@ -674,9 +683,22 @@
}
],
"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",
"pygments_lexer": "ipython3"
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.13"
}
},
"nbformat": 4,
Expand Down

0 comments on commit 59ea7a2

Please sign in to comment.