{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Forest Fire Model\n",
    "## A rapid introduction to Mesa\n",
    "\n",
    "The [Forest Fire Model](http://en.wikipedia.org/wiki/Forest-fire_model) is one of the simplest examples of a model that exhibits self-organized criticality.\n",
    "\n",
    "Mesa is a new, Pythonic agent-based modeling framework. A big advantage of using Python is that it a great language for interactive data analysis. Unlike some other ABM frameworks, with Mesa you can write a model, run it, and analyze it all in the same environment. (You don't have to, of course. But you can).\n",
    "\n",
    "In this notebook, we'll go over a rapid-fire (pun intended, sorry) introduction to building and analyzing a model with Mesa."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, some imports. We'll go over what all the Mesa ones mean just below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from mesa import Model, Agent\n",
    "from mesa.time import RandomActivation\n",
    "from mesa.space import Grid\n",
    "from mesa.datacollection import DataCollector\n",
    "from mesa.batchrunner import BatchRunner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the model\n",
    "\n",
    "Most models consist of basically two things: agents, and an world for the agents to be in. The Forest Fire model has only one kind of agent: a tree. A tree can either be unburned, on fire, or already burned. The environment is a grid, where each cell can either be empty or contain a tree.\n",
    "\n",
    "First, let's define our tree agent. The agent needs to be assigned **x** and **y** coordinates on the grid, and that's about it. We could assign agents a condition to be in, but for now let's have them all start as being 'Fine'. Since the agent doesn't move, and there is only at most one tree per cell, we can use a tuple of its coordinates as a unique identifier.\n",
    "\n",
    "Next, we define the agent's **step** method. This gets called whenever the agent needs to act in the world and takes the *model* object to which it belongs as an input. The tree's behavior is simple: If it is currently on fire, it spreads the fire to any trees above, below, to the left and the right of it that are not themselves burned out or on fire; then it burns itself out. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeCell(Agent):\n",
    "    '''\n",
    "    A tree cell.\n",
    "    \n",
    "    Attributes:\n",
    "        x, y: Grid coordinates\n",
    "        condition: Can be \"Fine\", \"On Fire\", or \"Burned Out\"\n",
    "        unique_id: (x,y) tuple. \n",
    "    \n",
    "    unique_id isn't strictly necessary here, but it's good practice to give one to each\n",
    "    agent anyway.\n",
    "    '''\n",
    "    def __init__(self, model, pos):\n",
    "        '''\n",
    "        Create a new tree.\n",
    "        Args:\n",
    "            pos: The tree's coordinates on the grid. Used as the unique_id\n",
    "        '''\n",
    "        super().__init__(pos, model)\n",
    "        self.pos = pos\n",
    "        self.unique_id = pos\n",
    "        self.condition = \"Fine\"\n",
    "        \n",
    "    def step(self):\n",
    "        '''\n",
    "        If the tree is on fire, spread it to fine trees nearby.\n",
    "        '''\n",
    "        if self.condition == \"On Fire\":\n",
    "            neighbors = self.model.grid.get_neighbors(self.pos, moore=False)\n",
    "            for neighbor in neighbors:\n",
    "                if neighbor.condition == \"Fine\":\n",
    "                    neighbor.condition = \"On Fire\"\n",
    "            self.condition = \"Burned Out\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to define the model object itself. The main thing the model needs is the grid, which the trees are placed on. But since the model is dynamic, it also needs to include time -- it needs a schedule, to manage the trees activation as they spread the fire from one to the other.\n",
    "\n",
    "The model also needs a few parameters: how large the grid is and what the density of trees on it will be. Density will be the key parameter we'll explore below.\n",
    "\n",
    "Finally, we'll give the model a data collector. This is a Mesa object which collects and stores data on the model as it runs for later analysis.\n",
    "\n",
    "The constructor needs to do a few things. It instantiates all the model-level variables and objects; it randomly places trees on the grid, based on the density parameter; and it starts the fire by setting all the trees on one edge of the grid (x=0) as being On \"Fire\".\n",
    "\n",
    "Next, the model needs a **step** method. Like at the agent level, this method defines what happens every step of the model. We want to activate all the trees, one at a time; then we run the data collector, to count how many trees are currently on fire, burned out, or still fine. If there are no trees left on fire, we stop the model by setting its **running** property to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ForestFire(Model):\n",
    "    '''\n",
    "    Simple Forest Fire model.\n",
    "    '''\n",
    "    def __init__(self, height, width, density):\n",
    "        '''\n",
    "        Create a new forest fire model.\n",
    "        \n",
    "        Args:\n",
    "            height, width: The size of the grid to model\n",
    "            density: What fraction of grid cells have a tree in them.\n",
    "        '''\n",
    "        # Initialize model parameters\n",
    "        self.height = height\n",
    "        self.width = width\n",
    "        self.density = density\n",
    "        \n",
    "        # Set up model objects\n",
    "        self.schedule = RandomActivation(self)\n",
    "        self.grid = Grid(height, width, torus=False)\n",
    "        self.dc = DataCollector({\"Fine\": lambda m: self.count_type(m, \"Fine\"),\n",
    "                                \"On Fire\": lambda m: self.count_type(m, \"On Fire\"),\n",
    "                                \"Burned Out\": lambda m: self.count_type(m, \"Burned Out\")})\n",
    "        \n",
    "        # Place a tree in each cell with Prob = density\n",
    "        for x in range(self.width):\n",
    "            for y in range(self.height):\n",
    "                if self.random.random() < self.density:\n",
    "                    # Create a tree\n",
    "                    new_tree = TreeCell(self, (x, y))\n",
    "                    # Set all trees in the first column on fire.\n",
    "                    if x == 0:\n",
    "                        new_tree.condition = \"On Fire\"\n",
    "                    self.grid[y][x] = new_tree\n",
    "                    self.schedule.add(new_tree)\n",
    "        self.running = True\n",
    "        \n",
    "    def step(self):\n",
    "        '''\n",
    "        Advance the model by one step.\n",
    "        '''\n",
    "        self.schedule.step()\n",
    "        self.dc.collect(self)\n",
    "        # Halt if no more fire\n",
    "        if self.count_type(self, \"On Fire\") == 0:\n",
    "            self.running = False\n",
    "    \n",
    "    @staticmethod\n",
    "    def count_type(model, tree_condition):\n",
    "        '''\n",
    "        Helper method to count trees in a given condition in a given model.\n",
    "        '''\n",
    "        count = 0\n",
    "        for tree in model.schedule.agents:\n",
    "            if tree.condition == tree_condition:\n",
    "                count += 1\n",
    "        return count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the model\n",
    "\n",
    "Let's create a model with a 100 x 100 grid, and a tree density of 0.6. Remember, ForestFire takes the arguments *height*, *width*, *density*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "fire = ForestFire(100, 100, 0.6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run the model until it's done (that is, until it sets its **running** property to False) just use the **run_model()** method. This is implemented in the Model parent object, so we didn't need to implement it above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "fire.run_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's all there is to it!\n",
    "\n",
    "But... so what? This code doesn't include a visualization, after all. \n",
    "\n",
    "Remember the data collector? Now we can put the data it collected into a pandas DataFrame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = fire.dc.get_model_vars_dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And chart it, to see the dynamics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.axes._subplots.AxesSubplot at 0x1819f757a58>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22712 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "results.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, the fire burned itself out after about 90 steps, with many trees left unburned. \n",
    "\n",
    "You can try changing the density parameter and rerunning the code above, to see how different densities yield different dynamics. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.axes._subplots.AxesSubplot at 0x1819fd724e0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24200 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fire = ForestFire(100, 100, 0.8)\n",
    "fire.run_model()\n",
    "results = fire.dc.get_model_vars_dataframe()\n",
    "results.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... But to really understand how the final outcome varies with density, we can't just tweak the parameter by hand over and over again. We need to do a batch run. \n",
    "\n",
    "## Batch runs\n",
    "\n",
    "Batch runs, also called parameter sweeps, allow use to systemically vary the density parameter, run the model, and check the output. Mesa provides a BatchRunner object which takes a model class, a dictionary of parameters and the range of values they can take and runs the model at each combination of these values. We can also give it reporters, which collect some data on the model at the end of each run and store it, associated with the parameters that produced it.\n",
    "\n",
    "For ease of typing and reading, we'll first create the parameters to vary and the reporter, and then assign them to a new BatchRunner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "fixed_params = dict(height=50, # Height and width are constant\n",
    "                 width=50)\n",
    "# Vary density from 0.01 to 1, in 0.01 increments:\n",
    "variable_params = dict(density=np.linspace(0,1,101)[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# At the end of each model run, calculate the fraction of trees which are Burned Out\n",
    "model_reporter = {\"BurnedOut\": lambda m: (ForestFire.count_type(m, \"Burned Out\") / \n",
    "                                          m.schedule.get_agent_count()) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the batch runner\n",
    "param_run = BatchRunner(ForestFire, variable_parameters=variable_params, \n",
    "                        fixed_parameters=fixed_params, model_reporters=model_reporter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the BatchRunner, which we've named param_run, is ready to go. To run the model at every combination of parameters (in this case, every density value), just use the **run_all()** method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100it [00:04, 11.23it/s]\n"
     ]
    }
   ],
   "source": [
    "param_run.run_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with the data collector, we can extract the data the batch runner collected into a dataframe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = param_run.get_model_vars_dataframe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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>density</th>\n",
       "      <th>Run</th>\n",
       "      <th>BurnedOut</th>\n",
       "      <th>height</th>\n",
       "      <th>width</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.01</td>\n",
       "      <td>0</td>\n",
       "      <td>0.025000</td>\n",
       "      <td>50</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>72</th>\n",
       "      <td>0.73</td>\n",
       "      <td>0</td>\n",
       "      <td>0.989983</td>\n",
       "      <td>50</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>71</th>\n",
       "      <td>0.72</td>\n",
       "      <td>0</td>\n",
       "      <td>0.992896</td>\n",
       "      <td>50</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>70</th>\n",
       "      <td>0.71</td>\n",
       "      <td>0</td>\n",
       "      <td>0.981069</td>\n",
       "      <td>50</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>69</th>\n",
       "      <td>0.70</td>\n",
       "      <td>0</td>\n",
       "      <td>0.980057</td>\n",
       "      <td>50</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    density  Run  BurnedOut  height  width\n",
       "0      0.01    0   0.025000      50     50\n",
       "72     0.73    0   0.989983      50     50\n",
       "71     0.72    0   0.992896      50     50\n",
       "70     0.71    0   0.981069      50     50\n",
       "69     0.70    0   0.980057      50     50"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, each row here is a run of the model, identified by its parameter values (and given a unique index by the Run column). To view how the BurnedOut fraction varies with density, we can easily just plot them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 1)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 8656 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(df.density, df.BurnedOut)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see the very clear emergence of a critical value around 0.5, where the model quickly shifts from almost no trees being burned, to almost all of them.\n",
    "\n",
    "In this case we ran the model only once at each value. However, it's easy to have the BatchRunner execute multiple runs at each parameter combination, in order to generate more statistically reliable results. We do this using the *iteration* argument.\n",
    "\n",
    "Let's run the model 5 times at each parameter point, and export and plot the results as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "500it [00:22, 11.33it/s] \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0, 1)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10832 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "param_run = BatchRunner(ForestFire, variable_params, fixed_params, \n",
    "                        iterations=5, model_reporters=model_reporter)\n",
    "param_run.run_all()\n",
    "df = param_run.get_model_vars_dataframe()\n",
    "plt.scatter(df.density, df.BurnedOut)\n",
    "plt.xlim(0,1)"
   ]
  }
 ],
 "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
}
