demo_transmission_tree.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb8dc11c-25c9-423c-899e-9c8fe5f14f90",
   "metadata": {},
   "source": [
    "# Demo: Transmission tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "9676957a-143f-444c-8501-60b31f52ab3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastsir import DiscreteSIR as sir\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e86e8471-c1cd-4c7e-8627-fba036196011",
   "metadata": {},
   "source": [
    "For the sake of simplicity, let's use simple contagion on WS networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "09034d89-20d3-494d-b818-085957c13c65",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Network parameters\n",
    "N = 10000\n",
    "m = 10\n",
    "p = 0.2\n",
    "G = nx.watts_strogatz_graph(N,m,p)\n",
    "max_degree = np.max([G.degree(n) for n in G])\n",
    "edgelist = list(G.edges())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "ecea235c-7c41-49d0-8df8-7ac52c377848",
   "metadata": {},
   "outputs": [],
   "source": [
    "#contagion parameters\n",
    "r = 0.01 #recovery probability\n",
    "q = 0.0015 #from the demo on phase transition, we know this is close to criticality\n",
    "recovery_probability = r\n",
    "infection_func = lambda k: 1-(1-q)**k\n",
    "infection_probability = infection_func(np.arange(max_degree+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "b2a0ad99-09e4-4245-8a0a-64854c9c5d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "#simulation parameters\n",
    "period = np.inf #simulate until no infected node\n",
    "#initialize object\n",
    "process = sir(edgelist,recovery_probability,infection_probability)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce58072-1f56-482f-82c3-4c2c95490f71",
   "metadata": {},
   "source": [
    "## Getting a single transmission tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "989c9379-52ca-4840-b4b3-3fce3098910c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#infect a random node\n",
    "initial_infected_nodes = {np.random.randint(N)}\n",
    "process.infect_node_set(initial_infected_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "4f88afe0-2bf0-4968-b08c-afec88bd8f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "#make the process evolve, mention we want to save the tree\n",
    "process.evolve(period,save_transmission_tree=True,save_macro_state=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "59883d0e-84e0-41aa-9024-24a2db20dfdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#get the transmission tree\n",
    "tree = process.get_transmission_tree()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c942ad7d-1d37-4ece-bd30-ede795030065",
   "metadata": {},
   "source": [
    "The \"transmission tree\" containes minimal information. It is a list of tuple, where each tuple include information about the transmission, in order: time of infection, generation of infector, infector id, infected id, infected degree of infectee (how many infected nodes where around the infectee at infection)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "265cb7db-0071-4458-a8a2-0aa2baedb6c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(15.0, 0, 1580, 1579, 1),\n",
       " (41.0, 1, 1579, 1577, 2),\n",
       " (151.0, 0, 4839, 4837, 1),\n",
       " (175.0, 1, 4837, 8346, 1),\n",
       " (189.0, 2, 8346, 8342, 1),\n",
       " (221.0, 0, 4839, 4834, 1),\n",
       " (225.0, 3, 8342, 8345, 2),\n",
       " (307.0, 4, 8345, 8348, 1),\n",
       " (401.0, 2, 8346, 8351, 1),\n",
       " (543.0, 5, 8348, 9270, 1),\n",
       " (546.0, 5, 8348, 5950, 1),\n",
       " (571.0, 5, 8348, 8349, 1),\n",
       " (606.0, 6, 9270, 9269, 1),\n",
       " (648.0, 6, 8349, 8347, 2),\n",
       " (652.0, 7, 9269, 9266, 1),\n",
       " (666.0, 6, 5950, 5953, 1),\n",
       " (719.0, 5, 8348, 8352, 2),\n",
       " (752.0, 6, 5950, 5954, 2),\n",
       " (774.0, 6, 8349, 8354, 1),\n",
       " (785.0, 6, 8352, 8357, 2),\n",
       " (811.0, 7, 8357, 8361, 1),\n",
       " (823.0, 7, 8357, 8353, 3),\n",
       " (827.0, 7, 5953, 5949, 2),\n",
       " (841.0, 7, 8354, 8359, 2),\n",
       " (846.0, 7, 8357, 8362, 2),\n",
       " (871.0, 6, 8352, 8350, 2),\n",
       " (880.0, 8, 8362, 8360, 1),\n",
       " (910.0, 8, 8353, 8355, 3),\n",
       " (1188.0, 7, 8350, 9074, 1)]"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96edb8ce-37b1-407b-8e64-3316bec6eb28",
   "metadata": {},
   "source": [
    "To parse the information in a pandas dataframe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "b7156a0e-8792-462c-b240-be9b9d4a6380",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>generation</th>\n",
       "      <th>infector</th>\n",
       "      <th>infectee</th>\n",
       "      <th>infected degree</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15</td>\n",
       "      <td>0</td>\n",
       "      <td>1580</td>\n",
       "      <td>1579</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>41</td>\n",
       "      <td>1</td>\n",
       "      <td>1579</td>\n",
       "      <td>1577</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>151</td>\n",
       "      <td>0</td>\n",
       "      <td>4839</td>\n",
       "      <td>4837</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>175</td>\n",
       "      <td>1</td>\n",
       "      <td>4837</td>\n",
       "      <td>8346</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>189</td>\n",
       "      <td>2</td>\n",
       "      <td>8346</td>\n",
       "      <td>8342</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>221</td>\n",
       "      <td>0</td>\n",
       "      <td>4839</td>\n",
       "      <td>4834</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>225</td>\n",
       "      <td>3</td>\n",
       "      <td>8342</td>\n",
       "      <td>8345</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>307</td>\n",
       "      <td>4</td>\n",
       "      <td>8345</td>\n",
       "      <td>8348</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>401</td>\n",
       "      <td>2</td>\n",
       "      <td>8346</td>\n",
       "      <td>8351</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>543</td>\n",
       "      <td>5</td>\n",
       "      <td>8348</td>\n",
       "      <td>9270</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>546</td>\n",
       "      <td>5</td>\n",
       "      <td>8348</td>\n",
       "      <td>5950</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>571</td>\n",
       "      <td>5</td>\n",
       "      <td>8348</td>\n",
       "      <td>8349</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>606</td>\n",
       "      <td>6</td>\n",
       "      <td>9270</td>\n",
       "      <td>9269</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>648</td>\n",
       "      <td>6</td>\n",
       "      <td>8349</td>\n",
       "      <td>8347</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>652</td>\n",
       "      <td>7</td>\n",
       "      <td>9269</td>\n",
       "      <td>9266</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>666</td>\n",
       "      <td>6</td>\n",
       "      <td>5950</td>\n",
       "      <td>5953</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>719</td>\n",
       "      <td>5</td>\n",
       "      <td>8348</td>\n",
       "      <td>8352</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>752</td>\n",
       "      <td>6</td>\n",
       "      <td>5950</td>\n",
       "      <td>5954</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>774</td>\n",
       "      <td>6</td>\n",
       "      <td>8349</td>\n",
       "      <td>8354</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>785</td>\n",
       "      <td>6</td>\n",
       "      <td>8352</td>\n",
       "      <td>8357</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>811</td>\n",
       "      <td>7</td>\n",
       "      <td>8357</td>\n",
       "      <td>8361</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>823</td>\n",
       "      <td>7</td>\n",
       "      <td>8357</td>\n",
       "      <td>8353</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>827</td>\n",
       "      <td>7</td>\n",
       "      <td>5953</td>\n",
       "      <td>5949</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>841</td>\n",
       "      <td>7</td>\n",
       "      <td>8354</td>\n",
       "      <td>8359</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>846</td>\n",
       "      <td>7</td>\n",
       "      <td>8357</td>\n",
       "      <td>8362</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>871</td>\n",
       "      <td>6</td>\n",
       "      <td>8352</td>\n",
       "      <td>8350</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>880</td>\n",
       "      <td>8</td>\n",
       "      <td>8362</td>\n",
       "      <td>8360</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>910</td>\n",
       "      <td>8</td>\n",
       "      <td>8353</td>\n",
       "      <td>8355</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>1188</td>\n",
       "      <td>7</td>\n",
       "      <td>8350</td>\n",
       "      <td>9074</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    time  generation  infector  infectee  infected degree\n",
       "0     15           0      1580      1579                1\n",
       "1     41           1      1579      1577                2\n",
       "2    151           0      4839      4837                1\n",
       "3    175           1      4837      8346                1\n",
       "4    189           2      8346      8342                1\n",
       "5    221           0      4839      4834                1\n",
       "6    225           3      8342      8345                2\n",
       "7    307           4      8345      8348                1\n",
       "8    401           2      8346      8351                1\n",
       "9    543           5      8348      9270                1\n",
       "10   546           5      8348      5950                1\n",
       "11   571           5      8348      8349                1\n",
       "12   606           6      9270      9269                1\n",
       "13   648           6      8349      8347                2\n",
       "14   652           7      9269      9266                1\n",
       "15   666           6      5950      5953                1\n",
       "16   719           5      8348      8352                2\n",
       "17   752           6      5950      5954                2\n",
       "18   774           6      8349      8354                1\n",
       "19   785           6      8352      8357                2\n",
       "20   811           7      8357      8361                1\n",
       "21   823           7      8357      8353                3\n",
       "22   827           7      5953      5949                2\n",
       "23   841           7      8354      8359                2\n",
       "24   846           7      8357      8362                2\n",
       "25   871           6      8352      8350                2\n",
       "26   880           8      8362      8360                1\n",
       "27   910           8      8353      8355                3\n",
       "28  1188           7      8350      9074                1"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree = np.array(tree, dtype=int)\n",
    "df = pd.DataFrame({'time':tree[:,0],'generation':tree[:,1],'infector':tree[:,2],'infectee':tree[:,3],'infected degree':tree[:,4]})\n",
    "df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}