1-sa_live_virus_on_network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sensitivity Analysis Made Easy with the EMA Workbench: Live coding with Raphael Klein"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Google Colab version*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the live-coding notebook of the workshop on *sensitivity analysis* (SA) at the Social Simulation Festival 2021. Here we will demonstrate how to do Variance-based SA also know as [Sobol SA](https://en.wikipedia.org/wiki/Variance-based_sensitivity_analysis) on a relatively simple model [virus on network](https://ccl.northwestern.edu/netlogo/models/VirusonaNetwork). The idea is that you reuse (read copy-paste) this code your own model. Therefore, we tried to keep simple and avoid extensive side steps from.\n",
    "\n",
    "This notebook is tuned to be run on [Google Colab](https://colab.research.google.com/) and has a couple of extra lines of code. If you want to use it on your local machine please use `sa_demo_local_machine.ipynb`.\n",
    "\n",
    "The core packages used in this notebook are [Mesa](https://mesa.readthedocs.io/en/stable/) to define an ABM model in Python, [EMA Workbench](https://emaworkbench.readthedocs.io/en/latest/) to design and run experiments, [SALib](https://salib.readthedocs.io/en/latest/) to conduct SA (within EMA Workbench). Also, we used one pretty plotting function of [pyNetLogo](https://pynetlogo.readthedocs.io/en/latest/).\n",
    "\n",
    "The notebook follows a simplified SA workflow and has 5 sections-steps:\n",
    "\n",
    "<img src=\"img/workflow.png\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Installations and imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the repo to make its file available for Google Colab\n",
    "!git clone https://github.com/BROSE-Uninc/SSF2021.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install necessary packages\n",
    "!pip install ema_workbench mesa ipyparallel SALib &> /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import datetime\n",
    "import random\n",
    "\n",
    "# Import EMA Workbench modules \n",
    "from ema_workbench import (ReplicatorModel, RealParameter, BooleanParameter, IntegerParameter, Constant, TimeSeriesOutcome, perform_experiments, save_results, ema_logging)\n",
    "\n",
    "# Initialize logger to keep track of experiments run\n",
    "ema_logging.log_to_stderr(ema_logging.INFO)\n",
    "\n",
    "# Import Mesa virus on network model\n",
    "from SSF2021.virus_on_network import model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The very first step of SA with EMA Workbench is to define or \"load\" the model as a function. That is, EMA Workbench treats all models as functions (read *black box*). They are supposed to have **inputs** (parameters, constants, uncertainties and policy levers) and **outputs** (outcomes, KPIs). The model structure is not interesting for EMA Workbench. It may be something simple as `def func(x)` which just returns x + 1.\n",
    "\n",
    "Our model has quite an extensive set of **8 inputs** model parameters (read more about the model [here](https://github.com/projectmesa/mesa/tree/master/examples/virus_on_network)). Their names are pretty self-explanatory, but let's quickly go through them:\n",
    "\n",
    "1. `num_nodes` - number of network nodes,\n",
    "2. `avg_node_degree` - average node degree,\n",
    "3. `initial_outbreak_size` - initial number of infected nodes,\n",
    "4. `virus_spread_chance` - chance of the spread of the virus,\n",
    "5. `virus_check_frequency` - how often node checks,\n",
    "6. `recovery_chance` - how likely the node recovers,\n",
    "7. `gain_resistance_chance` - what is the chance that node get resistance,\n",
    "8. `steps` - number of steps.\n",
    "\n",
    "And **4 outputs** model outcomes which correspond to a simple SIR model:\n",
    "1. `Susceptible` - number of agents in susceptible state, \n",
    "2. `Infected` - number of agents in infected state,\n",
    "3. `Resistant` - number of agents in resistant state,,\n",
    "4. `TIME` - to keep track of simulation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the model as a function\n",
    "def model_virus_on_network(num_nodes=1, \n",
    "                           avg_node_degree=1, \n",
    "                           initial_outbreak_size=1, \n",
    "                           virus_spread_chance=1, \n",
    "                           virus_check_frequency=1, \n",
    "                           recovery_chance=1, \n",
    "                           gain_resistance_chance=1,\n",
    "                           steps=10):\n",
    "    \n",
    "    from ... import ...\n",
    "    \n",
    "    # Initialising the model\n",
    "    virus_on_network = model.VirusOnNetwork(num_nodes=num_nodes, \n",
    "                                            avg_node_degree=avg_node_degree, \n",
    "                                            initial_outbreak_size=initial_outbreak_size, \n",
    "                                            virus_spread_chance=virus_spread_chance, \n",
    "                                            virus_check_frequency=virus_check_frequency, \n",
    "                                            recovery_chance=recovery_chance, \n",
    "                                            gain_resistance_chance=gain_resistance_chance)\n",
    "                \n",
    "    # Run the model steps times\n",
    "    virus_on_network.run_model(steps)\n",
    "\n",
    "    # Get model outcomes\n",
    "    outcomes = virus_on_network.datacollector.get_model_vars_dataframe()\n",
    "    \n",
    "    # Return model outcomes\n",
    "    return {'TIME' : list(range(steps + 1)),\n",
    "            \"...\" : outcomes[\"...\"].tolist(),\n",
    "            \"...\" : outcomes[\"...\"].tolist(),\n",
    "            \"...\" : outcomes[\"...\"].tolist()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's parameterize and test out our Mesa model. What is supposed to happen? First, we shouldn't get any error 😅. Second, after we run `model_virus_on_network` function it has to give us a set of model outcomes. Let's try."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parametrize the model\n",
    "num_nodes = ...\n",
    "avg_node_degree = ...\n",
    "initial_outbreak_size = ...\n",
    "virus_spread_chance = ...\n",
    "virus_check_frequency = ...\n",
    "recovery_chance = ...\n",
    "gain_resistance_chance = ...\n",
    "steps = ...\n",
    "\n",
    "model_virus_on_network(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great success 🥳! Let's visualize how the Mesa model outcomes will look like with this parameter set (for this we need to run `run.py` in ).\n",
    "\n",
    "<img src=\"img/run.png\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Design experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to design experiments. What does it mean? Well, we have to specify:\n",
    "\n",
    "* **which model parameters** aka *inputs* are we going to sample, what are their **ranges**, and random **distributions**,\n",
    "* what we will keep as **constants** and do not change over the model run,\n",
    "* and finally which **outcomes** we want to observe.\n",
    "\n",
    "It's an important step in SA workflow and we have to be careful. Because if parameter ranges are too narrow or they're sampled from e.g. a Normal distribution, there is a chance that you'll overlook import model behavior. This is why model parameters are named **uncertainties** in the EMA Workbench. We often do not know parameter vales and how to explore many plausible options."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's talk about \"tech\" part. First we have to initialize an instance of EMA Workbench called `ReplicatorModel`. This is how we \"connect\" EMA Workbench to our Python model. We have to pass a name of our model to `ReplicatorModel`, and also pass the function that we defined previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate and pass the model \n",
    "model = ReplicatorModel('...', function=...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Second, let's access the model and write down its attributes. Down below you can see that we have different classes for different types of parameters: `IntegerParameter`, `RealParameter`. We use the first one for integer parameters and the second for floats. To specify a parameter we have to pass it a name, left and right boundaries for sampling. As you can see, we don't to the same for `Constants`. There you have to specify only one value. Since our model it's not static and outcomes differ over time, we have to call `TimeSeriesOutcome`. Finally we need to specify how many replications do we want to have. Rule of thumb here, let's go with 10 😎."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model parameters and their ranges to be sampled\n",
    "#      Here include num_nodes, avg_node_degree, virus_spread_change, \n",
    "#      virus_check_frequency, recovery_chance and gain_resistance_chance\n",
    "model.uncertainties = [RealParameter(\"...\", ..., ...)]\n",
    "\n",
    "# Define model parameters that will remain constant\n",
    "#      \n",
    "model.constants = [Constant(\"initial_outbreak_size\", 1),\n",
    "                   Constant('....', ...)]\n",
    "\n",
    "# Define model outcomes\n",
    "model.outcomes = [TimeSeriesOutcome('TIME'),\n",
    "                  TimeSeriesOutcome('...')]\n",
    "\n",
    "# Define the number of replications\n",
    "model.replications = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Run the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we're all set to run the model. The syntax here is pretty straightforward: we call `perform_experiments`, passing it a model instance and specify how many `scenarios` we need. Depending on the number of scenarios EMA Workbench samples the parameters. As a result, the more scenarios you pass, the more parameter combinations you get. One more attribute here is the type of `uncertainty_sampling`. Some SA methods require a specific way of how parameters should be sampled, e.g. Sobol SA. That's why we have to specify it here.\n",
    "\n",
    "Important thing to remember: Sobol SA requires a relatively **large sample size**. As [docs](https://pynetlogo.readthedocs.io/en/latest/_docs/SALib_ipyparallel.html#Using-SALib-for-sensitivity-analysis) say, to calculate first-order, second-order and total sensitivity indices, we need to have sample size of n(2p+2), where p is the number of input parameters, and n is a baseline sample size which should be large enough to stabilize the estimation of the indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run experiments with the aforementioned parameters and outputs\n",
    "results = perform_experiments(models=..., scenarios=..., uncertainty_sampling='sobol')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the results\n",
    "experiments, outcomes = results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiments.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "outcomes.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Often it is a good idea to save the results of a lengthy run and open up it next day 🥱. Great that EMA Workbench has such a feature."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Epic save"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"img/save.jpg\" width=200 height=260/><center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ema_workbench.util.utilities import save_results, load_results\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creaet a directory to store the results\n",
    "directory = 'results/virus_on_network'\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the results\n",
    "save_results(results, directory + '/results.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the results\n",
    "results = load_results(directory + '/results.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiments, outcomes = results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A bit of preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we proceed further, we have to do a bit of preprocessing. Let's take a closer look at the outcomes. We have 4 outcomes, 1400 scenarios, 10 replications and we run the model for 30 steps. This is what you see down below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(random.choice(list(outcomes)))\n",
    "outcomes[random.choice(list(outcomes))].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simplify our life 😅, let's take a mean over all replications that we had. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_outcomes = {key:np.mean(outcomes[key],axis=1) for key in outcomes.keys()}\n",
    "mean_results = (experiments.copy(), mean_outcomes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now the shape of this array doesn't have 10 in it  \n",
    "mean_outcomes[random.choice(list(outcomes))].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visuals!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ema_workbench.analysis.plotting import lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.figsize'] = [10, 12]\n",
    "figure = lines(experiments, mean_outcomes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Do sensitivity analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we're here 🧙. EMA Workbench and SALib have different methods available for sampling and SA. Here we're going to use Variance-based Sensitivity Analysis also known as Sobol (by the surname of its author). Take a look at other options available there. As usual the most complex task here is to find out whether your model fits the method. As Confucius said \"*If you found a method that fits your problem within an hour, you're either a method developer or most probably bamboozled yourself.*\" \n",
    "For now, let's imagine that we found the right method and it is Sobol.\n",
    "\n",
    "The *tech* part here is relatively straightforward:\n",
    "\n",
    "1. Specify a *problem*, or simply say what model parameters did you sample,\n",
    "2. Select an *outcome* of interest, yes, we have to analyze the impact outcome by outcome,\n",
    "3. Run the analysis.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from SALib.analyze import sobol\n",
    "from ema_workbench.em_framework.salib_samplers import get_SALib_problem\n",
    "from SSF2021.src.plot import plot_sobol_indices\n",
    "sns.set_style('white')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the problem\n",
    "problem = get_SALib_problem(model.uncertainties)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select and normalize an outcome\n",
    "normalized_resistant = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform Sobol SA\n",
    "Si = sobol.analyze(problem=problem, Y=normalized_resistant,\n",
    "                   calc_second_order=True, print_to_console=False)\n",
    "\n",
    "# Get scores by type \n",
    "Si_filter = {k:Si[k] for k in ['ST', 'ST_conf', 'S1', 'S1_conf']}\n",
    "\n",
    "# Create a DataFrame out of them\n",
    "Si_df = pd.DataFrame(Si_filter, index=problem['names'])\n",
    "\n",
    "# Get indices and error bars\n",
    "indices = Si_df[['S1','ST']]\n",
    "err = Si_df[['S1_conf','ST_conf']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the results as a barplot\n",
    "fig, ax = plt.subplots(1)\n",
    "indices.plot.bar(yerr=err.values.T, ax=ax)\n",
    "fig.set_size_inches(8,6)\n",
    "fig.subplots_adjust(bottom=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Even a nicer plot\n",
    "fig = plot_sobol_indices(Si, problem, criterion='ST', threshold=0.005)\n",
    "fig.set_size_inches(7,7)"
   ]
  },
  {
   "source": [
    "## 5. Bonus - Interpretation of results"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "The box plot shows the contribution of each of the three parameters of interest to the output variance. The gain_resistance_chance parameter contributes to over 70% of the output variance (highest ST index). It's individual contribution is a little lower at 50%. One can note that for the recovery_chance parameter, the individual and total contributions are quite different with the first at 10% and the second at 50%. This gap is also present with the virus_check_frequency parameter but in a lower extent.\n",
    "\n",
    "The second result graph shows the interactions between the parameters. In the present case, one can observe that two parameters strongly interact: the recovery_chance and the virus_check_frequency parameters. This is symbolised by the strong grey link between them. Note that the circle show the same information as for the bar plot in the first results graph but in a different layout."
   ],
   "cell_type": "markdown",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}