{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introductory Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tutorial Description\n",
    "\n",
    "[Mesa](https://github.com/projectmesa/mesa) is a Python framework for [agent-based modeling](https://en.wikipedia.org/wiki/Agent-based_model).  Getting started with Mesa is easy. In this tutorial, we will walk through creating a simple model and progressively add functionality which will illustrate Mesa's core features.\n",
    "\n",
    "**Note:** This tutorial is a work-in-progress. If you find any errors or bugs, or just find something unclear or confusing, [let us know](https://github.com/projectmesa/mesa/issues)!\n",
    "\n",
    "The base for this tutorial is a very simple model of agents exchanging money. Next, we add *space* to allow agents move. Then, we'll cover two of Mesa's analytic tools: the *data collector* and *batch runner*. After that, we'll add an *interactive visualization* which lets us watch the model as it runs. Finally, we go over how to write your own visualization module, for users who are comfortable with JavaScript.\n",
    "\n",
    "You can also find all the code this tutorial describes in the **examples/Tutorial-Boltzmann_Wealth_Model** directory of the Mesa repository."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Model Description\n",
    "\n",
    "The tutorial model is a very simple simulated agent-based economy, drawn from econophysics and presenting a statistical mechanics approach to wealth distribution [Dragulescu2002]_. The rules of our tutorial model:\n",
    "\n",
    "1. There are some number of agents.\n",
    "2. All agents begin with 1 unit of money.\n",
    "3. At every step of the model, an agent gives 1 unit of money (if they have it) to some other agent.\n",
    "\n",
    "Despite its simplicity, this model yields results that are often unexpected to those not familiar with it. For our purposes, it also easily demonstrates Mesa's core features.\n",
    "\n",
    "Let's get started."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation\n",
    "\n",
    "To start, install Mesa. We recommend doing this in a [virtual environment](https://virtualenvwrapper.readthedocs.org/en/stable/), but make sure your environment is set up with Python 3. Mesa requires Python3 and does not work in Python 2 environments.\n",
    "\n",
    "To install Mesa, simply:\n",
    "\n",
    "```bash\n",
    "    $ pip install mesa\n",
    "```\n",
    "\n",
    "When you do that, it will install Mesa itself, as well as any dependencies that aren't in your setup yet. Additional dependencies required by this tutorial can be found in the **examples/Tutorial-Boltzmann_Wealth_Model/requirements.txt** file, which can be installed by running:\n",
    "\n",
    "```bash\n",
    "    $ pip install -r requirements.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a sample model\n",
    "\n",
    "Once Mesa is installed, you can start building our model. You can write models in two different ways:\n",
    "\n",
    "1. Write the code in its own file with your favorite text editor, or\n",
    "2. Write the model interactively in [Jupyter Notebook](http://jupyter.org/) cells.\n",
    "\n",
    "Either way, it's good practice to put your model in its own folder -- especially if the project will end up consisting of multiple files (for example, Python files for the model and the visualization, a Notebook for analysis, and a Readme with some documentation and discussion).\n",
    "\n",
    "Begin by creating a folder, and either launch a Notebook or create a new Python source file.  We will use the name `MoneyModel.py` here.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up the model\n",
    "\n",
    "To begin writing the model code, we start with two core classes: one for the overall model, the other for the agents. The model class holds the model-level attributes, manages the agents, and generally handles the global level of our model. Each instantiation of the model class will be a specific model run. Each model will contain multiple agents, all of which are instantiations of the agent class. Both the model and agent classes are child classes of Mesa's generic `Model` and `Agent` classes.\n",
    "\n",
    "Each agent has only one variable: how much wealth it currently has. (Each agent will also have a unique identifier (i.e., a name), stored in the `unique_id` variable. Giving each agent a unique id is a good practice when doing agent-based modeling.)\n",
    "\n",
    "There is only one model-level parameter: how many agents the model contains. When a new model is started, we want it to populate itself with the given number of agents.\n",
    "\n",
    "The beginning of both classes looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesa import Agent, Model\n",
    "\n",
    "class MoneyAgent(Agent):\n",
    "    \"\"\"An agent with fixed initial wealth.\"\"\"\n",
    "    def __init__(self, unique_id, model):\n",
    "        super().__init__(unique_id, model)\n",
    "        self.wealth = 1\n",
    "\n",
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N):\n",
    "        self.num_agents = N\n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding the scheduler\n",
    "\n",
    "Time in most agent-based models moves in steps, sometimes also called **ticks**. At each step of the model, one or more of the agents -- usually all of them -- are activated and take their own step, changing internally and/or interacting with one another or the environment.\n",
    "\n",
    "The **scheduler** is a special model component which controls the order in which agents are activated. For example, all the agents may activate in the same order every step; their order might be shuffled; we may try to simulate all the agents acting at the same time; and more. Mesa offers a few different built-in scheduler classes, with a common interface. That makes it easy to change the activation regime a given model uses, and see whether it changes the model behavior. This may not seem important, but scheduling patterns can have an impact on your results [Comer2014]_.\n",
    "\n",
    "For now, let's use one of the simplest ones: `RandomActivation`, which activates all the agents once per step, in random order. Every agent is expected to have a ``step`` method, which takes a model object as its only argument -- this is the agent's action when it is activated. We add an agent to the schedule using the `add` method; when we call the schedule's `step` method, it shuffles the order of the agents, then activates them all, one at a time.\n",
    "\n",
    "With that in mind, the model code with the scheduler added looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesa import Agent, Model\n",
    "from mesa.time import RandomActivation\n",
    "\n",
    "class MoneyAgent(Agent):\n",
    "    \"\"\" An agent with fixed initial wealth.\"\"\"\n",
    "    def __init__(self, unique_id, model):\n",
    "        super().__init__(unique_id, model)\n",
    "        self.wealth = 1\n",
    "\n",
    "    def step(self):\n",
    "        # The agent's step will go here.\n",
    "        pass\n",
    "\n",
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N):\n",
    "        self.num_agents = N\n",
    "        self.schedule = RandomActivation(self)\n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)\n",
    "            self.schedule.add(a)\n",
    "\n",
    "    def step(self):\n",
    "        '''Advance the model by one step.'''\n",
    "        self.schedule.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, we have a model which runs -- it just doesn't do anything. You can see for yourself with a few easy lines. If you've been working in an interactive session, you can create a model object directly. Otherwise, you need to open an interactive session in the same directory as your source code file, and import the classes. For example, if your code is in `MoneyModel.py`:\n",
    "\n",
    "```python\n",
    "from MoneyModel import MoneyModel\n",
    "```\n",
    "Then create the model object, and run it for one step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "empty_model = MoneyModel(10)\n",
    "empty_model.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise\n",
    "    \n",
    "Try modifying the code above to have every agent print out its `unique_id` when it is activated. Run a few steps of the model to see how the agent activation order is shuffled each step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Agent Step\n",
    "\n",
    "Now we just need to have the agents do what we intend for them to do: check their wealth, and if they have the money, give one unit of it away to another random agent. To allow the agent to choose another agent at random, we use the ``model.random`` random-number generator. This works just like Python's ``random`` module, but with a fixed seed set when the model is instantiated, that can be used to replicate a specific model run later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To pick an agent at random, we need a list of all agents. Notice that there isn't such a list explicitly in the model. The scheduler, however, does have an internal list of all the agents it is scheduled to activate.\n",
    "\n",
    "With that in mind, we rewrite the agent's `step` method, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoneyAgent(Agent):\n",
    "    \"\"\" An agent with fixed initial wealth.\"\"\"\n",
    "    def __init__(self, unique_id, model):\n",
    "        super().__init__(unique_id, model)\n",
    "        self.wealth = 1\n",
    "\n",
    "    def step(self):\n",
    "        if self.wealth == 0:\n",
    "            return\n",
    "        other_agent = self.random.choice(self.model.schedule.agents)\n",
    "        other_agent.wealth += 1\n",
    "        self.wealth -= 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running your first model\n",
    "\n",
    "With that last piece in hand, it's time for the first rudimentary run of the model.\n",
    "\n",
    "If you've written the code in its own file (`MoneyModel.py` or a different name), launch an interpreter in the same directory as the file (either the plain Python command-line interpreter, or the IPython interpreter), or launch a Jupyter Notebook there. Then import the classes you created. (If you wrote the code in a Notebook, obviously this step isn't necessary).\n",
    "\n",
    "```python\n",
    "from MoneyModel import *\n",
    "```\n",
    "\n",
    "Now let's create a model with 10 agents, and run it for 10 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MoneyModel(10)\n",
    "for i in range(10):\n",
    "    model.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to get some data out of the model. Specifically, we want to see the distribution of the agent's wealth. We can get the wealth values with list comprehension, and then use matplotlib (or another graphics library) to visualize the data in a histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([5., 0., 0., 2., 0., 0., 1., 0., 0., 2.]),\n",
       " array([0. , 0.3, 0.6, 0.9, 1.2, 1.5, 1.8, 2.1, 2.4, 2.7, 3. ]),\n",
       " <a list of 10 Patch objects>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 4188 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# For a jupyter notebook add the following line:\n",
    "%matplotlib inline\n",
    "\n",
    "# The below is needed for both notebooks and scripts\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "agent_wealth = [a.wealth for a in model.schedule.agents]\n",
    "plt.hist(agent_wealth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are running from a text editor or IDE, you'll also need to add this line, to make the graph appear.\n",
    "\n",
    "```python\n",
    "plt.show()\n",
    "```\n",
    "\n",
    "You'll probably see something like the distribution shown below. Yours will almost certainly look at least slightly different, since each run of the model is random, after all.\n",
    "\n",
    "To get a better idea of how a model behaves, we can create multiple model runs and see the distribution that emerges from all of them. We can do this with a nested for loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([447., 284., 162.,  59.,  29.,  19.]),\n",
       " array([0, 1, 2, 3, 4, 5, 6]),\n",
       " <a list of 6 Patch objects>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 4732 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_wealth = []\n",
    "for j in range(100):\n",
    "    # Run the model\n",
    "    model = MoneyModel(10)\n",
    "    for i in range(10):\n",
    "        model.step()\n",
    "    \n",
    "    # Store the results\n",
    "    for agent in model.schedule.agents:\n",
    "        all_wealth.append(agent.wealth)\n",
    "\n",
    "plt.hist(all_wealth, bins=range(max(all_wealth)+1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This runs 100 instantiations of the model, and runs each for 10 steps. (Notice that we set the histogram bins to be integers, since agents can only have whole numbers of wealth). This distribution looks a lot smoother. By running the model 100 times, we smooth out some of the 'noise' of randomness, and get to the model's overall expected behavior.\n",
    "\n",
    "This outcome might be surprising. Despite the fact that all agents, on average, give and receive one unit of money every step, the model converges to a state where most agents have a small amount of money and a small number have a lot of money."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding space\n",
    "\n",
    "Many ABMs have a spatial element, with agents moving around and interacting with nearby neighbors. Mesa currently supports two overall kinds of spaces: grid, and continuous. Grids are divided into cells, and agents can only be on a particular cell, like pieces on a chess board. Continuous space, in contrast, allows agents to have any arbitrary position. Both grids and continuous spaces are frequently [toroidal](https://en.wikipedia.org/wiki/Toroidal_graph), meaning that the edges wrap around, with cells on the right edge connected to those on the left edge, and the top to the bottom. This prevents some cells having fewer neighbors than others, or agents being able to go off the edge of the environment.\n",
    "\n",
    "Let's add a simple spatial element to our model by putting our agents on a grid and make them walk around at random. Instead of giving their unit of money to any random agent, they'll give it to an agent on the same cell.\n",
    "\n",
    "Mesa has two main types of grids: `SingleGrid` and `MultiGrid`. `SingleGrid` enforces at most one agent per cell; `MultiGrid` allows multiple agents to be in the same cell. Since we want agents to be able to share a cell, we use `MultiGrid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesa.space import MultiGrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We instantiate a grid with width and height parameters, and a boolean as to whether the grid is toroidal. Let's make width and height model parameters, in addition to the number of agents, and have the grid always be toroidal. We can place agents on a grid with the grid's `place_agent` method, which takes an agent and an (x, y) tuple of the coordinates to place the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N, width, height):\n",
    "        self.num_agents = N\n",
    "        self.grid = MultiGrid(width, height, True)\n",
    "        self.schedule = RandomActivation(self)\n",
    "        \n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)\n",
    "            self.schedule.add(a)\n",
    "            \n",
    "            # Add the agent to a random grid cell\n",
    "            x = self.random.randrange(self.grid.width)\n",
    "            y = self.random.randrange(self.grid.height)\n",
    "            self.grid.place_agent(a, (x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Under the hood, each agent's position is stored in two ways: the agent is contained in the grid in the cell it is currently in, and the agent has a `pos` variable with an (x, y) coordinate tuple. The `place_agent` method adds the coordinate to the agent automatically.\n",
    "\n",
    "Now we need to add to the agents' behaviors, letting them move around and only give money to other agents in the same cell.\n",
    "\n",
    "First let's handle movement, and have the agents move to a neighboring cell. The grid object provides a `move_agent` method, which like you'd imagine, moves an agent to a given cell. That still leaves us to get the possible neighboring cells to move to. There are a couple ways to do this. One is to use the current coordinates, and loop over all coordinates +/- 1 away from it. For example:\n",
    "\n",
    "```python\n",
    "neighbors = []\n",
    "x, y = self.pos\n",
    "for dx in [-1, 0, 1]:\n",
    "    for dy in [-1, 0, 1]:\n",
    "        neighbors.append((x+dx, y+dy))\n",
    "```\n",
    "\n",
    "But there's an even simpler way, using the grid's built-in `get_neighborhood` method, which returns all the neighbors of a given cell. This method can get two types of cell neighborhoods: Moore (including diagonals), and Von Neumann (only up/down/left/right). It also needs an argument as to whether to include the center cell itself as one of the neighbors.\n",
    "\n",
    "With that in mind, the agent's `move` method looks like this:\n",
    "\n",
    "```python\n",
    "class MoneyAgent(Agent):\n",
    "   #...\n",
    "    def move(self):\n",
    "        possible_steps = self.model.grid.get_neighborhood(\n",
    "            self.pos, \n",
    "            moore=True,\n",
    "            include_center=False)\n",
    "        new_position = self.random.choice(possible_steps)\n",
    "        self.model.grid.move_agent(self, new_position)\n",
    "```\n",
    "\n",
    "Next, we need to get all the other agents present in a cell, and give one of them some money. We can get the contents of one or more cells using the grid's `get_cell_list_contents` method, or by accessing a cell directly. The method accepts a list of cell coordinate tuples, or a single tuple if we only care about one cell.\n",
    "\n",
    "```python\n",
    "class MoneyAgent(Agent):\n",
    "    #...\n",
    "    def give_money(self):\n",
    "        cellmates = self.model.grid.get_cell_list_contents([self.pos])\n",
    "        if len(cellmates) > 1:\n",
    "            other = self.random.choice(cellmates)\n",
    "            other.wealth += 1\n",
    "            self.wealth -= 1\n",
    "```\n",
    "\n",
    "And with those two methods, the agent's ``step`` method becomes:\n",
    "\n",
    "```python\n",
    "class MoneyAgent(Agent):\n",
    "    # ...\n",
    "    def step(self):\n",
    "        self.move()\n",
    "        if self.wealth > 0:\n",
    "            self.give_money()\n",
    "```\n",
    "\n",
    "Now, putting that all together should look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N, width, height):\n",
    "        self.num_agents = N\n",
    "        self.grid = MultiGrid(width, height, True)\n",
    "        self.schedule = RandomActivation(self)\n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)\n",
    "            self.schedule.add(a)\n",
    "            # Add the agent to a random grid cell\n",
    "            x = self.random.randrange(self.grid.width)\n",
    "            y = self.random.randrange(self.grid.height)\n",
    "            self.grid.place_agent(a, (x, y))\n",
    "\n",
    "    def step(self):\n",
    "        self.schedule.step()\n",
    "\n",
    "class MoneyAgent(Agent):\n",
    "    \"\"\" An agent with fixed initial wealth.\"\"\"\n",
    "    def __init__(self, unique_id, model):\n",
    "        super().__init__(unique_id, model)\n",
    "        self.wealth = 1\n",
    "\n",
    "    def move(self):\n",
    "        possible_steps = self.model.grid.get_neighborhood(\n",
    "            self.pos, \n",
    "            moore=True, \n",
    "            include_center=False)\n",
    "        new_position = self.random.choice(possible_steps)\n",
    "        self.model.grid.move_agent(self, new_position)\n",
    "\n",
    "    def give_money(self):\n",
    "        cellmates = self.model.grid.get_cell_list_contents([self.pos])\n",
    "        if len(cellmates) > 1:\n",
    "            other = self.random.choice(cellmates)\n",
    "            other.wealth += 1\n",
    "            self.wealth -= 1\n",
    "\n",
    "    def step(self):\n",
    "        self.move()\n",
    "        if self.wealth > 0:\n",
    "            self.give_money()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a model with 50 agents on a 10x10 grid, and run it for 20 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MoneyModel(50, 10, 10)\n",
    "for i in range(20):\n",
    "    model.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's use matplotlib and numpy to visualize the number of agents residing in each cell. To do that, we create a numpy array of the same size as the grid, filled with zeros. Then we use the grid object's `coord_iter()` feature, which lets us loop over every cell in the grid, giving us each cell's coordinates and contents in turn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x2612cbb19e8>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 8288 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "agent_counts = np.zeros((model.grid.width, model.grid.height))\n",
    "for cell in model.grid.coord_iter():\n",
    "    cell_content, x, y = cell\n",
    "    agent_count = len(cell_content)\n",
    "    agent_counts[x][y] = agent_count\n",
    "plt.imshow(agent_counts, interpolation='nearest')\n",
    "plt.colorbar()\n",
    "\n",
    "# If running from a text editor or IDE, remember you'll need the following:\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Collecting Data\n",
    "\n",
    "So far, at the end of every model run, we've had to go and write our own code to get the data out of the model. This has two problems: it isn't very efficient, and it only gives us end results. If we wanted to know the wealth of each agent at each step, we'd have to add that to the loop of executing steps, and figure out some way to store the data.\n",
    "\n",
    "Since one of the main goals of agent-based modeling is generating data for analysis, Mesa provides a class which can handle data collection and storage for us and make it easier to analyze.\n",
    "\n",
    "The data collector stores three categories of data: model-level variables, agent-level variables, and tables (which are a catch-all for everything else). Model- and agent-level variables are added to the data collector along with a function for collecting them. Model-level collection functions take a model object as an input, while agent-level collection functions take an agent object as an input. Both then return a value computed from the model or each agent at their current state. When the data collector’s `collect` method is called, with a model object as its argument, it applies each model-level collection function to the model, and stores the results in a dictionary, associating the current value with the current step of the model. Similarly, the method applies each agent-level collection function to each agent currently in the schedule, associating the resulting value with the step of the model, and the agent’s `unique_id`.\n",
    "\n",
    "Let's add a DataCollector to the model, and collect two variables. At the agent level, we want to collect every agent's wealth at every step. At the model level, let's measure the model's [Gini Coefficient](https://en.wikipedia.org/wiki/Gini_coefficient), a measure of wealth inequality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesa.datacollection import DataCollector\n",
    "\n",
    "def compute_gini(model):\n",
    "    agent_wealths = [agent.wealth for agent in model.schedule.agents]\n",
    "    x = sorted(agent_wealths)\n",
    "    N = model.num_agents\n",
    "    B = sum( xi * (N-i) for i,xi in enumerate(x) ) / (N*sum(x))\n",
    "    return (1 + (1/N) - 2*B)\n",
    "\n",
    "class MoneyAgent(Agent):\n",
    "    \"\"\" An agent with fixed initial wealth.\"\"\"\n",
    "    def __init__(self, unique_id, model):\n",
    "        super().__init__(unique_id, model)\n",
    "        self.wealth = 1\n",
    "\n",
    "    def move(self):\n",
    "        possible_steps = self.model.grid.get_neighborhood(\n",
    "            self.pos, \n",
    "            moore=True, \n",
    "            include_center=False)\n",
    "        new_position = self.random.choice(possible_steps)\n",
    "        self.model.grid.move_agent(self, new_position)\n",
    "\n",
    "    def give_money(self):\n",
    "        cellmates = self.model.grid.get_cell_list_contents([self.pos])\n",
    "        if len(cellmates) > 1:\n",
    "            other = self.random.choice(cellmates)\n",
    "            other.wealth += 1\n",
    "            self.wealth -= 1\n",
    "\n",
    "    def step(self):\n",
    "        self.move()\n",
    "        if self.wealth > 0:\n",
    "            self.give_money()\n",
    "\n",
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N, width, height):\n",
    "        self.num_agents = N\n",
    "        self.grid = MultiGrid(width, height, True)\n",
    "        self.schedule = RandomActivation(self)\n",
    "        \n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)\n",
    "            self.schedule.add(a)\n",
    "            # Add the agent to a random grid cell\n",
    "            x = self.random.randrange(self.grid.width)\n",
    "            y = self.random.randrange(self.grid.height)\n",
    "            self.grid.place_agent(a, (x, y))\n",
    "        \n",
    "        self.datacollector = DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini},\n",
    "            agent_reporters={\"Wealth\": \"wealth\"})\n",
    "\n",
    "    def step(self):\n",
    "        self.datacollector.collect(self)\n",
    "        self.schedule.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At every step of the model, the datacollector will collect and store the model-level current Gini coefficient, as well as each agent's wealth, associating each with the current step.\n",
    "\n",
    "We run the model just as we did above. Now is when an interactive session, especially via a Notebook, comes in handy: the DataCollector can export the data it's collected as a pandas DataFrame, for easy interactive analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MoneyModel(50, 10, 10)\n",
    "for i in range(100):\n",
    "    model.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the series of Gini coefficients as a pandas DataFrame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.axes._subplots.AxesSubplot at 0x2612cbcf048>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16636 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gini = model.datacollector.get_model_vars_dataframe()\n",
    "gini.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we can get the agent-wealth data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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></th>\n",
       "      <th>Wealth</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Step</th>\n",
       "      <th>AgentID</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"5\" valign=\"top\">0</th>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              Wealth\n",
       "Step AgentID        \n",
       "0    0             1\n",
       "     1             1\n",
       "     2             1\n",
       "     3             1\n",
       "     4             1"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_wealth = model.datacollector.get_agent_vars_dataframe()\n",
    "agent_wealth.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll see that the DataFrame's index is pairings of model step and agent ID. You can analyze it the way you would any other DataFrame. For example, to get a histogram of agent wealth at the model's end:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.axes._subplots.AxesSubplot at 0x2612d972470>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 7516 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "end_wealth = agent_wealth.xs(99, level=\"Step\")[\"Wealth\"]\n",
    "end_wealth.hist(bins=range(agent_wealth.Wealth.max()+1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or to plot the wealth of a given agent (in this example, agent 14):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.axes._subplots.AxesSubplot at 0x2612dacec50>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16032 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "one_agent_wealth = agent_wealth.xs(14, level=\"AgentID\")\n",
    "one_agent_wealth.Wealth.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batch Run\n",
    "\n",
    "Like we mentioned above, you usually won't run a model only once, but multiple times, with fixed parameters to find the overall distributions the model generates, and with varying parameters to analyze how they drive the model's outputs and behaviors. Instead of needing to write nested for-loops for each model, Mesa provides a BatchRunner class which automates it for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesa.batchrunner import BatchRunner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The BatchRunner also requires an additional variable running for the MoneyModel class. This variable enables conditional shut off of the model once a condition is met. In this example it will be set as True indefinitely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoneyModel(Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "    def __init__(self, N, width, height):\n",
    "        self.num_agents = N\n",
    "        self.grid = MultiGrid(width, height, True)\n",
    "        self.schedule = RandomActivation(self)\n",
    "        self.running = True\n",
    "        \n",
    "        # Create agents\n",
    "        for i in range(self.num_agents):\n",
    "            a = MoneyAgent(i, self)\n",
    "            self.schedule.add(a)\n",
    "            # Add the agent to a random grid cell\n",
    "            x = self.random.randrange(self.grid.width)\n",
    "            y = self.random.randrange(self.grid.height)\n",
    "            self.grid.place_agent(a, (x, y))\n",
    "        \n",
    "        self.datacollector = DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini},\n",
    "            agent_reporters={\"Wealth\": \"wealth\"})\n",
    "\n",
    "    def step(self):\n",
    "        self.datacollector.collect(self)\n",
    "        self.schedule.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We instantiate a BatchRunner with a model class to run, and two dictionaries: one of the fixed parameters (mapping model arguments to values) and one of varying parameters (mapping each parameter name to a sequence of values for it to take).  The BatchRunner also takes an argument for how many model instantiations to create and run at each combination of parameter values, and how many steps to run each instantiation for. Finally, like the DataCollector, it takes dictionaries of model- and agent-level reporters to collect. Unlike the DataCollector, it won't collect the data every step of the model, but only at the end of each run.\n",
    "\n",
    "In the following example, we hold the height and width fixed, and vary the number of agents. We tell the BatchRunner to run 5 instantiations of the model with each number of agents, and to run each for 100 steps. We have it collect the final Gini coefficient value.\n",
    "\n",
    "Now, we can set up and run the BatchRunner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "245it [01:30,  1.41it/s]\n"
     ]
    }
   ],
   "source": [
    "fixed_params = {\"width\": 10,\n",
    "               \"height\": 10}\n",
    "variable_params = {\"N\": range(10, 500, 10)}\n",
    "\n",
    "batch_run = BatchRunner(MoneyModel, \n",
    "                        variable_params,\n",
    "                        fixed_params,\n",
    "                        iterations=5, \n",
    "                        max_steps=100,\n",
    "                        model_reporters={\"Gini\": compute_gini})\n",
    "batch_run.run_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like the DataCollector, we can extract the data we collected as a DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x261302b8f98>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 14188 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run_data = batch_run.get_model_vars_dataframe()\n",
    "run_data.head()\n",
    "plt.scatter(run_data.N, run_data.Gini)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that each row is a model run, and gives us the parameter values associated with that run. We can use  this data to view a scatter-plot comparing the number of agents to the final Gini."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Happy Modeling!\n",
    "\n",
    "This document is a work in progress.  If you see any errors, exclusions or have any problems please contact [us](https://github.com/projectmesa/mesa/issues)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`virtual environment`: http://docs.python-guide.org/en/latest/dev/virtualenvs/\n",
    "\n",
    "[Comer2014] Comer, Kenneth W. “Who Goes First? An Examination of the Impact of Activation on Outcome Behavior in AgentBased Models.” George Mason University, 2014. http://gradworks.umi.com/36/23/3623940.html.\n",
    "\n",
    "[Dragulescu2002] Drăgulescu, Adrian A., and Victor M. Yakovenko. “Statistical Mechanics of Money, Income, and Wealth: A Short Survey.” arXiv Preprint Cond-mat/0211175, 2002. http://arxiv.org/abs/cond-mat/0211175."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:mesa]",
   "language": "python",
   "name": "conda-env-mesa-py"
  },
  "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.6.6"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
