ABC of ABM in Julia

TL;DR: When writing agent-based models in Julia try to use a single agent type to get good performance. If you definitely need more than one type of agent you still can get a good performance but it requires a bit more complex design of your code.

Introduction

In this post I discuss basic approaches to implementing Agent Based Models (ABM) in Julia. It covers a fragment of a tutorial that I will be giving with Przemysław Szufel at Social Simulation Conference 2018, workshop Running  high performance simulations with Julia programming language on Monday, August 20.

The post summarizes some thoughts about issues raised in recent discussion on Discourse about Agent Based Modeling in Julia.

While there are many possible approaches to implementation of ABMs in Julia I want to concentrate on basic techniques that can be picked up by someone who just starts learning Julia.

Our working example is implementation of forest fire model which is described in detail an excellent book An Introduction to Agent-Based Modeling by Uri Wilensky and William Rand. We will exactly reproduce the NetLogo implementation model. In particular I will avoid certain possible optimizations of the code to keep the organization of the logic follow NetLogo implementation.

This post is divided into three sections:

1. Explaining how NetLogo model works
2. Implementation in Julia using a single type of agent
3. Implementation in Julia using several types of agents

All examples in this post should be run under Julia 0.7, currently in beta. I will update the codes if something would start to break after Julia 0.7 is released and that is why in this post they are linked as gists (here is a link for the impatient).

Also I assume that you know Julia a bit (during the workshop at the conference all will be explained starting from the basics).

Forest fire model

We have a 251x251 rectangular grid. Initially each cell of the grid is empty or contains a tree.

A tree has three possible states: green, on fire and burnt. Initially all trees are green and a tree is present in a cell with probability density which is a parameter of the model.

Now how the model works. In the initial step we set that all cells in the first row of the grid to contain trees that are on fire. Next in each step:

1. we select all trees that are on fire;
2. we iterate through them in a random order;
3. for each tree on fire if it touches a green tree then the green tree is set on fire;
4. finally the on fire tree changes state to burnt.

Our question is what percentage of trees that are initially present will get burnt in the process.

As a reference on my laptop running 100 replications of this model for density=0.55 takes around 30 seconds with all animations and updating disabled, and for density=0.75 it is over one minute (I do not try here or below to do very precise benchmarks as I want to concentrate on orders of magnitude).

A single type of agent

In this model implementing it with a single type of agent (a tree) is natural. Such an implementation can be expected to be easily made efficient in Julia. The reason is that we will have all containers (vectors, matrices, sets, dictionaries, etc.) hold a single concrete type. The benefit of this is the following in terms of performance:

1. Julia compiler should be able to infer types of all variables in all functions (we know that we have only one type of agent).
2. In particular (and this is often crucial) Julia compiler knows what method of a function it should dispatch if some method has the agent as a parameter (e.g. action of the agent).

In this case the power of Julia is that mostly, when you think about performance, you do not care if the type to represent an agent is in-built into Julia or your custom type nor whether it is an immutable or mutable type (there are differences and probably there are cases when they are significant but I want to stress the first level of thinking).

To see this consider two implementations of the model. The first one uses Int (in-built, immutable) to represent agent state on the grid, the second uses custom Tree type (user-defined, mutable and storing some more information than Int-version):

1. Version using integers: forestfire1.jl
2. Version using custom type: forestfire2.jl

As you can see the model with Tree type does a bit more work but essentially the code logic is very similar. Here are timings of running the codes:

$ julia7 forestfire1.jl

  2.190497 seconds (1.11 M allocations: 112.423 MiB, 0.94% gc time)

  5.586829 seconds (465.93 k allocations: 303.833 MiB, 0.92% gc time)

$ julia7 forestfire2.jl

  2.924750 seconds (7.44 M allocations: 305.734 MiB, 3.42% gc time)

  7.770683 seconds (6.76 M allocations: 495.277 MiB, 6.90% gc time)

The version using integers is a bit faster as expected but they are both significantly (10x) faster than NetLogo and the timings are of the same order of magnitude.

Several types of agents

In this example using one type of agent is natural, but let us test what happens if we force several types of agents into the model. Specifically we notice that in forestfire2.jl we have when field meaningful only for burned (brown) tree. So we decide to use three separate types of agents TreeGreen, TreeRed and TreeBrown. Additionally then we denote cell without a tree with nothing.

The implementation of such a model is given in file forestfire3.jl. The problem with it is that it is much slower as grid matrix has type Any (you can test yourself that making all tree types a subtype of some abstract type or making type of the matrix a union does not change what we get below). Therefore we can expect that it will be much slower. This is confirmed by running the model:

$ julia7 forestfire3.jl

 37.078864 seconds (694.45 M allocations: 20.809 GiB, 4.40% gc time)

 90.306497 seconds (2.04 G allocations: 60.961 GiB, 5.49% gc time)

and we see that we are roughly at the speed level of NetLogo.

The good thing is that Julia allows us to write such a code and in many cases it will be fast enough. In particular the code is much slower because agents to a lot of very simple actions so the cost of iteration is much larger than the cost of actions themselves. If agents had a complex and expensive logic then it could be moved out to a function (a technique called barrier functions) and the overhead of type instability would not be that significant.

However, the question is if we can make code fast using agents of heterogeneous types. Here we will consider the simplest possible technique that allows to achieve this. What you essentially do is:

1. store information about agents in a tuple, I call it trees in the code
2. each entry of this tuple is a collection of agents of a single type (in the example we will use a vector but the choice of collection should be tailored to the needs of the simulation)
3. you create a single type, I call it TreeID in the code that allows you to select an appropriate element from the tuple in a type stable way; in our example it holds two fields:
• typ identifying agent type (number of slot within a tuple)
• loc identifying agent location (position of agent within the collection that is a slot of a tuple)
4. the crucial thing is that the trees tuple holding collections of agents of homogeneous type should be always indexed by a number known at compile time (alternatively you could create a struct and select its fields) - this ensures that all usages of trees tuple will allow the compiler to infer the type of the result (in short what you have to avoid is passing a variable to index trees tuple; the general pattern is to use a sequence of if-elseif-elseif... statements based on the value of typ in TreeID)

The code implementing this pattern is given here forestfire4.jl. I have even complicated it a bit on purpose by adding x and y fields to TreeRed and defining burn function that has to be called to show that Julia is able to handle them at compile time. The downside is that the code got a bit more complex. We have the following mapping of typ value in TreeID:

• 0 means no tree (thus no mapping to trees is needed)
• 1 means green tree
• 2 means red tree
• 3 means brown tree

The crucial question is what is the performance of this pattern. Here is a result of running the code:

$ julia7 forestfire4.jl

  2.403108 seconds (1.04 M allocations: 171.372 MiB, 1.24% gc time)

  6.376856 seconds (505.42 k allocations: 653.171 MiB, 2.38% gc time)

And we see that it is very good.

The crucial benefits of this pattern are the following (by ID-structure I call an equivalent of TreeID in a general code):

1. You can iterate over agents in whatever order you want (the ID-structure does not force you to process types of agents in separate batches)
2. You can perform actions that rely on type of agent without having to reach to the agent; you can do it on ID-structure level;
3. You can use ID-structure anywhere you want (it can be in action scheduler, it can be in a representation of locations of agents in space, it can be in a graph of connections, ...)
4. If you have methods that should have different implementations depending on agent type then passing them ID-structure and using if-elseif-elseif... template inside you can store the logic that depends on the tuple-container (or struct-container) structure only in a few places of your code and most of the time not have to care about it by working on ID-structure level.