agentpy_workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AgentPy Workshop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p>\n",
    "<a href=\"https://mybinder.org/v2/gh/JoelForamitti/agentpy_workshop/HEAD\" rel=\"noopener\" target=\"_blank\", style=\"float: left; padding-right:4px\"><img src=\"https://mybinder.org/badge_logo.svg\" alt=\"Binder\"></a>\n",
    "<a href=\"https://agentpy.readthedocs.io/en/latest/\" rel=\"noopener\" target=\"_blank\"><img src=\"https://img.shields.io/badge/read_the-documentation-red?logo=read-the-docs\" alt=\"Documentation\", style=\"float: left; padding-right:4px\"></a> \n",
    "<a href=\"https://img.shields.io/github/license/JoelForamitti/agentpy_workshop\" rel=\"noopener\" target=\"_blank\"><img src=\"https://img.shields.io/badge/edit_on-GitHub-blueviolet?logo=github\" alt=\"GitHub\", style=\"float: left; padding-right:4px\"></a>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This interactive notebook will guide you through the development of an agent-based model with the [AgentPy](https://agentpy.readthedocs.io) package (v0.1.5). \n",
    "To run this workshop in a virtual environment, you can launch this notebook in [binder]([![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/JoelForamitti/agentpy_workshop/HEAD))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Introduction\n",
    "\n",
    "**Agent-based models** (ABMs) are computer simulations to generate and understand the dynamics of [complex systems](https://complexityexplained.github.io/) based on the autonomous behavior of heterogeneous agents. They show how the behavior of a system as a whole can emerge out of the actions and interactions of individual agents, each with their own set of rules, characteristics, and information. This method has found numerous applications in fields like [ecology](https://doi.org/10.3389/fevo.2018.00237), [cognitive sciences](https://doi.org/10.1177/0963721419834547), [management](https://oxford.universitypressscholarship.com/view/10.1093/acprof:oso/9780195172119.001.0001/acprof-9780195172119), \n",
    "[economics](https://www.nature.com/articles/460685a), and [sociology](https://doi.org/10.1002/wics.1356).\n",
    "\n",
    "**AgentPy** is an open-source library for the development and analysis of agent-based models in Python. The framework integrates the tasks of model design, numerical experiments, and data analysis within a single environment, and is optimized for interactive computing with [IPython](http://ipython.org/) and [Jupyter](https://jupyter.org/). If you have questions or ideas for improvements, please visit the [discussion forum](https://github.com/JoelForamitti/agentpy/discussions).\n",
    "\n",
    "Within this framework, **agents** are represented as [objects](https://www.programiz.com/python-programming/class) that can be placed within different **environments** like a network or a spatial grid. A **model** is used to initiate these objects, perform a simulation, and record data. **Experiments** can take a model and run it over multiple iterations with different parameter combinations. The resulting **output data** can be saved, loaded, and re-arranged for analysis and visualization.\n",
    "\n",
    "<img src=\"structure.png\" alt=\"drawing\" width=\"500\"/>\n",
    "\n",
    "**This workshop** will introduce you to the basic syntax of the agentpy package. Some basic knowledge of Python is required, but it is generally aimed at beginners. The following sections will show you how to design custom agents and models; use networks; visualize the dynamics of a simulation in an animation, perform experiments that run a model multiple times with different parameters; and run a sensitivity analysis on the output data. For more information about agentpy, check out the [documentation](https://agentpy.readthedocs.io/) and [model library](https://agentpy.readthedocs.io/en/latest/model_library.html). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Basic Syntax\n",
    "### Let's get started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below you can see the first *code cell* of this notebook, which you can edit yourself. Here are some basic functionalities of Jupyter:\n",
    "\n",
    "- You can navigate between cells with your keyboard arrows or by clicking on a cell. \n",
    "- You can switch between *Command Mode* the *Edit Mode* with `Enter` and `Esc`. \n",
    "- In Command Mode, you can create a new cell below with the command `b` and delete a cell with `d` `d`. \n",
    "- To *run* a selected cell, click on the `►` button at the top or use the keyboard combination `Shift` + `Enter`. \n",
    "- If the [kernel](https://ipython.org/ipython-doc/3/development/kernels.html) crashes, you can restart it with the `↻` button. \n",
    "\n",
    "For more information about how to use Jupyter, [click here](https://realpython.com/jupyter-notebook-introduction/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Hello World!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this workshop, we have to install the required packages `agentpy` and `seaborn`. On binder, this has already happened automatically. If you are running this workshop on your local machine, you can use the following command (remove the # to use the command):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -r requirements.txt "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now import the agentpy library as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import agentpy as ap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating custom agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us first understand the basics of object-oriented programming. **Objects** are a collection of variables and methods. Every object has a type, which refers to the class that has been used to create the object. **Classes** can be seen as a *building instruction* for new objects. You can learn more about classes and objects [here](https://www.programiz.com/python-programming/class). \n",
    "\n",
    "To create an agent-based model in agentpy, we create a new instance of the class [`ap.Model`](https://agentpy.readthedocs.io/en/latest/reference_model.html):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ap.Model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model is currently empty. To fill it with agents, we need to create a custom class as a *child* of [`Agent`](https://agentpy.readthedocs.io/en/latest/reference_agents.html). This means that our custom agent will inherit all basic agent functionalities from its *parent* class. You can find the description of the [`Agent`](https://agentpy.readthedocs.io/en/latest/reference_agents.html) class and it's methods in the [API reference](https://agentpy.readthedocs.io/en/latest/reference_agents.html).\n",
    "\n",
    "The following code creates a custom class named `MyAgentType`. This class defines a method [`setup`](https://agentpy.readthedocs.io/en/latest/reference_agents.html#agentpy.Agent.setup) that will be called automatically when a new agent of this class is created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyAgentType(ap.Agent):\n",
    "    \"\"\" Description of my agent. \"\"\"\n",
    "    \n",
    "    def setup(self):\n",
    "        # Called automatically when a new agent of this type is created\n",
    "        self.x = 0  # Initiate a variable\n",
    "        self.say_hello()  # Call a method\n",
    "        \n",
    "    def say_hello(self):\n",
    "        # Can be called multiple times throughout the simulation\n",
    "        print(f\"Agent {self.id} says 'Hello World!'\")\n",
    "        return self.id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To add three agents of this class to our model, we can use the following method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agents = ap.AgentList(model, 3, MyAgentType)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Agents are held within an [`AgentList`](https://agentpy.readthedocs.io/en/latest/reference_sequences.html#agentpy.AgentList), which allow for easy operations on all agents at the same time. Here are some examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Accessing attributes\n",
    "print(agents)\n",
    "print(agents.id)\n",
    "print(agents.type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manipulating variables\n",
    "agents.x = 0 \n",
    "print(agents.x)\n",
    "agents.x += 1\n",
    "print(agents.x)\n",
    "agents.x += agents.id\n",
    "print(agents.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calling methods\n",
    "return_values = agents.say_hello()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "return_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Selecting agents (by condition)\n",
    "selection = agents.select(agents.x != 3)\n",
    "print(selection)\n",
    "print(selection.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Selecting agents (randomly)\n",
    "random_selection = agents.random(2)\n",
    "print(random_selection)\n",
    "print(random_selection.x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating custom models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now proceed to define a model capable of performing a simple simulation. For this, we have to create a custom model class, similar to how we created a custom agent. This model comes with two methods: \n",
    " `Model.setup()`, which is called at the beginning of the simulation; and `Model.step()`, which will be called at every time-step of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModelType(ap.Model):\n",
    "    \"\"\" Description of my model. \"\"\"\n",
    "    \n",
    "    def setup(self):\n",
    "        # Called automatically at the beginning of a simulation\n",
    "        self.agents = ap.AgentList(self, self.p.agents, MyAgentType)  # Add new agents\n",
    "\n",
    "    def step(self):\n",
    "        # Called automatically at each time-step of the simulation\n",
    "        self.agents.x += 1  # Manipulate a variable\n",
    "        self.agents.record('x')  # Record a variable     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, this model will initiate some agents of our custom type and give them a variable `x` at the start of the simulation. Then, at every time-step, this variable will be increased by 1. The function [`Agent.record()`](https://agentpy.readthedocs.io/en/latest/overview.html#recording-data) is used to record this variable for later analysis.\n",
    "\n",
    "The number of agents and steps is given by a dictionary of **parameters**, which have to be passed to the model when a new instance is created, and can be accessed by all objects with `self.p`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {'agents':3, 'steps':3}\n",
    "model = MyModelType(parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use the function [`model.run()`](https://agentpy.readthedocs.io/en/latest/overview.html#running-a-simulation) to perform the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation results\n",
    "We have successfully performed our first simulation!\n",
    "We can now find a [`DataDict`](https://agentpy.readthedocs.io/en/latest/overview.html#output-and-analysis) with output data of this simulation in our variable `results`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Different sections of the results can be accessed like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.variables.MyAgentType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data structure of this table is a **pandas dataframe**, which come with numerous functionalities for data manipulation and analysis. You can learn more about the pandas library [here](https://pandas.pydata.org/docs/getting_started/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "Feel free experiment a bit with the code above. Here are some things you can try out:\n",
    "\n",
    "- Add another method to MyAgentType and use it in `MyModelType.step()`.\n",
    "- Add another parameter to the simulation and use it in the new method.\n",
    "- Add another variable to `MyAgentType` and record it as well.\n",
    "- Analyze the output data with the function [`pandas.DataFrame.mean()`](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.mean.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Creating Complex Models\n",
    "\n",
    "Let us now move on to a more advanced exercise. \n",
    "We want to create an agent-based model that simulates the propagation of a disease through a network.\n",
    "The agents of this model are people, which can be in one of the following three conditions: susceptible to the disease (S), infected (I), or recovered (R). The agents are connected to each other through a network of peers. At every time-step, infected agents can infect their peers or recover from the disease based on random chance. \n",
    "\n",
    "For the generation of the network graph, we will use the [Watts-Strogatz model](https://en.wikipedia.org/wiki/Watts%E2%80%93Strogatz_model). This is an algorithm that generates a regular network were every agent will have the same amount of connections, and then introduces a certain amount of randomness to change some of these connections. A network where most agents are not neighbors, but where it is easy to reach every other agent in a small number of steps, is called a [small-world network](https://en.wikipedia.org/wiki/Small-world_network).\n",
    "\n",
    "<img src=\"networks.png\" alt=\"drawing\" width=\"600\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, we import all the necessary libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model design\n",
    "import agentpy as ap\n",
    "import networkx as nx \n",
    "\n",
    "# Visualization\n",
    "import matplotlib.pyplot as plt \n",
    "import seaborn as sns\n",
    "import IPython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Next, we define a new agent type `Person`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Person(ap.Agent):\n",
    "    \"\"\" A person that can transmit a disease. \"\"\"\n",
    "    \n",
    "    def setup(self):  \n",
    "        \"\"\" Initialize a new variable at agent creation. \"\"\"\n",
    "        self.condition = 0  # Susceptible = 0, Infected = 1, Recovered = 2\n",
    "        \n",
    "    def being_sick(self):\n",
    "        \"\"\" Spread disease to peers in the network. \"\"\"\n",
    "        rg = self.model.random  # Random number generator\n",
    "        for n in self.network.neighbors(self): \n",
    "            if n.condition == 0 and self.p.infection_chance > rg.random():\n",
    "                n.condition = 1  # Infect susceptible peer\n",
    "        if self.p.recovery_chance > rg.random():\n",
    "            self.condition = 2  # Recover from infection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two new functions are used within this class:\n",
    "\n",
    "- `Agent.neighbors()` returns a list of the agents' peers in the network\n",
    "- `Model.random.random()` returns a uniform random draw between 0 and 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Finally, we define a new model type `VirusModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VirusModel(ap.Model):\n",
    "    \"\"\" Agent-based model that simulates the \n",
    "    propagation of a disease through a network. \"\"\"\n",
    "    \n",
    "    def setup(self):\n",
    "        \"\"\" Initialize the agents and network of the model. \"\"\"\n",
    "        \n",
    "        # Prepare a small-world network\n",
    "        graph = nx.watts_strogatz_graph(\n",
    "            self.p.population, \n",
    "            self.p.number_of_neighbors, \n",
    "            self.p.network_randomness)\n",
    "        \n",
    "        # Create agents and network\n",
    "        self.agents = ap.AgentList(self, self.p.population, Person)\n",
    "        self.network = self.agents.network = ap.Network(self, graph)\n",
    "        self.network.add_agents(self.agents, self.network.nodes)\n",
    "\n",
    "        # Infect a random share of the population\n",
    "        I0 = int(self.p.initial_infection_share * self.p.population)\n",
    "        self.agents.random(I0).condition = 1 \n",
    "\n",
    "    def update(self):  \n",
    "        \"\"\" Record variables after setup and each step. \"\"\"\n",
    "        \n",
    "        # Record share of agents with each condition\n",
    "        for i, c in enumerate(('S', 'I', 'R')):\n",
    "            n_agents = len(self.agents.select(self.agents.condition == i))\n",
    "            self[c] = n_agents / self.p.population \n",
    "            self.record(c)\n",
    "        \n",
    "        # Stop simulation if disease is gone\n",
    "        if self.I == 0:\n",
    "            self.stop()\n",
    "        \n",
    "    def step(self):   \n",
    "        \"\"\" Define the models' events per simulation step. \"\"\"\n",
    "        \n",
    "        # Call 'being_sick' for infected agents\n",
    "        self.agents.select(self.agents.condition == 1).being_sick()\n",
    "    \n",
    "    def end(self):     \n",
    "        \"\"\" Record evaluation measures at the end of the simulation. \"\"\"\n",
    "        \n",
    "        # Record final evaluation measures\n",
    "        self.report('Total share infected', self.I + self.R) \n",
    "        self.report('Peak share infected', max(self.log['I']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to `setup` and `step`, two new kinds of methods are defined in this model:\n",
    "\n",
    "- `Model.update()` is called automatically after both `setup` and `step`. This method is meant for the recording of variables, as it will also record variables at `t==0` (before the first step).\n",
    "- `Model.end()` is called at the end of the simulation. This method is meant to evaluate the simulation outcomes and record [evaluation measure](https://agentpy.readthedocs.io/en/latest/overview.html#recording-data).\n",
    "\n",
    "Some other methods and attributes of `Model` that are used:\n",
    "\n",
    "- `Model.add_network()` creates a new [network](https://agentpy.readthedocs.io/en/latest/reference_environments.html#agentpy.Network) environment.\n",
    "- `Model.measure()` records an [evaluation measure](https://agentpy.readthedocs.io/en/latest/overview.html#recording-data).\n",
    "- `Model.stop()` is used to stop the simulation.\n",
    "- `Model.log` accesses recorded variables.\n",
    "\n",
    "For a full overview of the methods and attributes of `Model`, check out the [API reference](https://agentpy.readthedocs.io/en/latest/reference_models.html).\n",
    "\n",
    "Another function that is used is [`nx.watts_strogatz_graph`](https://networkx.org/documentation/networkx-1.10/reference/generators.html), which is an imported function from the networkx library.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running a simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "As before, we can now define our parameters and run the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {  \n",
    "    'population': 1000,\n",
    "    'infection_chance': 0.3,\n",
    "    'recovery_chance': 0.1,\n",
    "    'initial_infection_share': 0.1,\n",
    "    'number_of_neighbors': 2,\n",
    "    'network_randomness': 0.5  \n",
    "}\n",
    "\n",
    "model = VirusModel(parameters)\n",
    "results = model.run() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyzing results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "To visualize the evolution of our variables over time, we create a stackplot function using the [matplotlib](https://matplotlib.org/) library. To understand the syntax of matplotlib, you can check out the documentation for the stackplot function [here](https://matplotlib.org/stable/gallery/lines_bars_and_markers/stackplot_demo.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def virus_stackplot(data, ax):\n",
    "    \"\"\" Stackplot of people's condition over time. \"\"\"\n",
    "    x = data.index.get_level_values('t')\n",
    "    y = [data[var] for var in ['I', 'S', 'R']]\n",
    "\n",
    "    sns.set()\n",
    "    ax.stackplot(x, y, labels=['Infected', 'Susceptible', 'Recovered'],\n",
    "                 colors = ['r', 'b', 'g'])\n",
    "\n",
    "    ax.legend()\n",
    "    ax.set_xlim(0, max(1, len(x)-1))\n",
    "    ax.set_ylim(0, 1)\n",
    "    ax.set_xlabel(\"Time steps\")\n",
    "    ax.set_ylabel(\"Percentage of population\")\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "virus_stackplot(results.variables.VirusModel, ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating an animation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "We can also animate the model's dynamics as follows.\n",
    "The function `animation_plot` takes a model instance \n",
    "and displays the previous stackplot together with a network graph. \n",
    "The function [`animate`](https://agentpy.readthedocs.io/en/latest/reference_analysis.html#agentpy.animate) will call this plot\n",
    "function for every time-step and return an matplotlib animation object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def animation_plot(m, axs):\n",
    "    ax1, ax2 = axs\n",
    "    ax1.set_title(\"Virus spread\")\n",
    "    ax2.set_title(f\"Share infected: {m.I}\")\n",
    "\n",
    "    # Plot stackplot on first axis\n",
    "    virus_stackplot(m.output.variables.VirusModel, ax1)\n",
    "\n",
    "    # Plot network on second axis\n",
    "    color_dict = {0:'b', 1:'r', 2:'g'}\n",
    "    colors = [color_dict[c] for c in m.agents.condition]\n",
    "    nx.draw_circular(m.network.graph, node_color=colors,\n",
    "                     node_size=50, ax=ax2)\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(8, 4)) # Prepare figure\n",
    "parameters['population'] = 50 # Lower population for better visibility\n",
    "animation = ap.animate(VirusModel(parameters), fig, axs, animation_plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using Jupyter, we can display this animation directly in our notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IPython.display.HTML(animation.to_jshtml()) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-run experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To explore the effect of different parameter values, we use the classes `Sample`, `Range`, and `IntRange` to create a sample of different parameter combinations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    'population': ap.IntRange(100, 1000),\n",
    "    'infection_chance': ap.Range(0.1, 1.),\n",
    "    'recovery_chance': ap.Range(0.1, 1.),\n",
    "    'initial_infection_share': 0.1,\n",
    "    'number_of_neighbors': 2,\n",
    "    'network_randomness': ap.Range(0., 1.)\n",
    "}\n",
    "\n",
    "sample = ap.Sample(\n",
    "    parameters,\n",
    "    n=128,\n",
    "    method='saltelli',\n",
    "    calc_second_order=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "We then create an `Experiment` that takes a model and sample as input. [`Experiment.run()`](https://agentpy.readthedocs.io/en/latest/reference_experiments.html#agentpy.Experiment.run) runs our model repeatedly over the whole sample with ten random iterations per parameter combination."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp = ap.Experiment(VirusModel, sample, iterations=10)\n",
    "results = exp.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally, we can save and load our results as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = ap.DataDict.load('VirusModel')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "The measures in our `DataDict` now hold one row for each simulation run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "We can use standard functions of the pandas library like \n",
    "`pandas.DataFrame.hist` to look at summary statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.reporters.hist()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sensitivity analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "[Sensitivity analysis](https://en.wikipedia.org/wiki/Sensitivity_analysis) can be used to understand how variations in a model's output relate to variations in the model's input parameters. Here, we calculate [Sobol Sensitivity Indices](https://salib.readthedocs.io/en/latest/api.html#sobol-sensitivity-analysis), which show us how much of the variance of each measure is caused by each varied parameter.\n",
    "The function [`sensitivity_sobol`](https://agentpy.readthedocs.io/en/latest/reference_analysis.html#agentpy.sensitivity_sobol) uses the \n",
    "[SAlib](https://salib.readthedocs.io/en/latest/basics.html) package to calculate these indices\n",
    "for the passed results and parameter ranges. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.calc_sobol()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "This adds two new categories to our results:\n",
    "\n",
    "- `sobol` returns first-order sobol sensitivity indices\n",
    "- `sobol_conf` returns confidence ranges for the above indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use pandas to create a bar plot that visualizes these sensitivity indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sobol(results):\n",
    "    \"\"\" Bar plot of Sobol sensitivity indices. \"\"\"\n",
    "\n",
    "    sns.set()\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8, 4))\n",
    "    si_list = results.sensitivity.sobol.groupby(by='reporter')\n",
    "    si_conf_list = results.sensitivity.sobol_conf.groupby(by='reporter')\n",
    "\n",
    "    for (key, si), (_, err), ax in zip(si_list, si_conf_list, axs):\n",
    "        si = si.droplevel('reporter')\n",
    "        err = err.droplevel('reporter')\n",
    "        si.plot.barh(xerr=err, title=key, ax=ax, capsize = 3)\n",
    "        ax.set_xlim(0)\n",
    "\n",
    "    axs[0].get_legend().remove()\n",
    "    axs[1].set(ylabel=None, yticklabels=[])\n",
    "    axs[1].tick_params(left=False)\n",
    "    plt.tight_layout()\n",
    "\n",
    "plot_sobol(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another way to visualize sensitivities is to plot the\n",
    "average evaluation measures over our parameter variations. \n",
    "For this, we can use the [seaborn](https://seaborn.pydata.org/) for statistical data visualization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sensitivity(results):\n",
    "    \"\"\" Show average simulation results for different parameter values. \"\"\"\n",
    "\n",
    "    sns.set()\n",
    "    fig, axs = plt.subplots(2, 2, figsize=(8, 8))\n",
    "    axs = [i for j in axs for i in j] # Flatten list\n",
    "\n",
    "    data = results.arrange_reporters().astype('float')\n",
    "    params = results.parameters.sample.keys()\n",
    "\n",
    "    for x, ax in zip(params, axs):\n",
    "        for y in results.reporters.columns:\n",
    "            sns.regplot(x=x, y=y, data=data, ax=ax, ci=99,\n",
    "                        x_bins=15, fit_reg=False, label=y)\n",
    "        ax.set_ylim(0, 1)\n",
    "        ax.set_ylabel('')\n",
    "        ax.legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "\n",
    "plot_sensitivity(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "If you want to experiment with this model further, here are some things you can try to do:\n",
    "\n",
    "- Try out different kinds of [network generators](https://networkx.org/documentation/networkx-1.10/reference/generators.html)\n",
    "- Implement a different infection process\n",
    "- Allow agents to be re-infected after they recovered\n",
    "- Implement a fourth condition: Vaccinated\n",
    "- Implement a vaccination process at each time-step\n",
    "- Make the vaccination speed dependent on a parameter\n",
    "- Analyze the sensitivity of this parameter "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Conclusion\n",
    "\n",
    "We have reached the end of this basic workshop on how to use the agentpy package. I hope you enjoyed it, and that everything was clear. If you have any feedback for improvements, please write to joel.foramitti@uab.cat. If you are interested to contribute, please check out this [guide](https://agentpy.readthedocs.io/en/latest/contributing.html).\n",
    "\n",
    "Now you are ready to start building your own models! For inspiration, check out the `/notebooks` folder in this repository, which contains various guides and demonstration models. To read about all features, please visit the [documentation](https://agentpy.readthedocs.io/). And of course feel free to ask questions and share your models in the [discussion forum](https://github.com/JoelForamitti/agentpy/discussions)."
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}