Mesa Code Segments

The mesa library is very nice, and it should be faster and easier to implement models that a plain numpy implementation.

However, mesa is undergoing changes so many of the online examples are based on earlier versions of the library. This page is a collection of code segments that I used to build all examples covered in class and what I suggest you use for assignments. These are based on the latest version of the library (3.0.0b0).

Using a consistent style between models/applications means that it is easier to focus on the similarities between the applications and should result in faster development.

As we implement more complicated models and need more functionality I will add to this page.
However, this page is not intended to be a replacement for the mesa documentation so spend some time browsing that.

Setup

First load in the standard modules — I have broken this into three cells based on their general purpose

1
2
3
4
5
6
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("darkgrid")

import pandas as pd

Next load functions to build animations 

1
2
from matplotlib import animation
from IPython.display import HTML

Finally install mesa using 

1
!pip install -U --pre mesa

and load in mesa and verify the correct version is used

1
2
import mesa
mesa.__version__

You should get output '3.0.2' or '3.1.0.dev'.

Agent class

We need to create a class to represent the Agent. Depending on the application, I might called this class Cell (for Game of Life model), Tree (for fire model), Person (for SIR models), etc., or I might be lazy and just call them Agent. In this page of code snippets I will always use name Agent.

The minimal code for Agent is

1
2
3
4
5
6
class Agent(mesa.Agent):

    def __init__(self, model):
        super().__init__(model)

    def step(self): ...

Agent class with states

Often an agent can be in one of a number of predefined states, so take the Game of Life where the agent, Cell, is either DEAD or ALIVE. In this case our class definition starts with:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Agent(mesa.Agent):

    DEAD = 0
    ALIVE = 1

    def __init__(self, model, state):
        super().__init__(model)
        self.state = state

    def step(self): ...

Agent class with simultaneous update

In some applications, such as the Game of Life, we want to update ALL agents at exactly the same time — in contrast to updating agents either in order (for the Traffic model) or at random.

To get the simultaneous update in mesa we have class

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Agent(mesa.Agent):

    DEAD = 0
    ALIVE = 1

    def __init__(self, model, state=DEAD):
        super().__init__(model)
        self.state = state
        self._next_state = state

    def compute_step(self): ...

    def step(self):
        self.state = self._next_state

In this code, our agent has:

Then in the step function in the Model class we have 

1
2
3
4
5
6
7
class Model(mesa.Model):

    ...

    def step(self):
        self.agents.do("compute_step")
        self.agents.do("step")

where the first line updates the _next_state for all agents without changing state and once that is completed the second line updates the state for all agents.

Model class

The Model class stores the agents, controls how that are updated, and what data is recorded. Similar to the Agent class you can call this anything, so you could use FireModel (for fire model), SchellingModel (for the Schelling model), etc., but I tend to just use Model. In this page of code snippets I will always use name Model.

At a minimum the Model class has

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Model(mesa.Model):

    def __init__(self, seed=None):
        super().__init__(seed=seed)

        # code to save model parameters

        # code to create world/space/grid

        # code to add agents

        # code to start simulation


    def step(self): ...
        # code to update model

Visualisation

Mesa comes with a nice visualisation and UI (for model parameter selection) but I had issues in getting it to work with Colab so we are going build our own using matplotllib.animation. We need to implement the following:

1
2
3
4
5
6
7
8
9
def draw_model(model, ax):

    # setup matplotlib window

    # draw the model grid (if model has space/grid)

    # draw agents in model based on their state

    # set title

1
2
3
def frame(n):
    model.step()
    draw_model(model, ax)

Then to create an animation we:

1
2
n_frames = 50
fig, ax = plt.subplots(nrows=1, ncols=1)

1
2
3
anim = animation.FuncAnimation(fig, frame, frames=n_frames)
plt.close()
HTML(anim.to_jshtml())

DataCollector class

See mesa documentation for details.

Miscellaneous Tasks

Loop over all locations in a grid

If Model has a grid, we often need to loop over all locations in the grid, for example to place an agent at each location (think Cell in Game of Life or a Tree in Fire Model). This can be done from within a method of Model using

1
2
for contents, pos in self.grid.coord_iter():
    pass

or if not interested in the contents (say when in Model.__init__ and placing agents) we use

1
2
3
for _, pos in self.grid.coord_iter():
    agent = Agent(model)
    self.grid.place_agent(agent, pos)

Loop over agents in model

To loop over all agents in the model use

1
2
for agent in self.agents:
    pass

List all neighbours of an agent

Within the Agent class, we often need to get a list of its neighbours — and then filter/count those alive, etc. The cleanest approach is to have a method which performs the task of getting the list of neighbours. For example, in the Game of Life we might have

1
2
3
4
5
6
class Agent(mesa.Agent):

    ...

    def get_neighbors(self):
        return self.model.grid.get_neighbors(self.pos, moore=True, include_center=False)