Introduction
In a previous article, I provide an introductory example (optimizing the number of supermarket counters) of how one can run agent based models using Python’s Mesa library. We tracked line charts of KPIs such as as the average queue length and waiting time of our customers, but we used limited visualization.
In this article, we’ll go through a different agent-based modeling example where we’ll focus more on visualization.
Note from Towards Data Science’s editors: While we allow independent authors to publish articles in accordance with our rules and guidelines, we do not endorse each author’s contribution. You should not rely on an author’s works without seeking professional advice. See our Reader Terms for details.
COVID spread agent-based model
We’ll model how COVID spreads using a model inspired by the SIR model (susceptible, infected, recovered). We have a 50 x 50 grid with 250 cells and a number of agents (people) that move between adjacent cells, potentially infecting one another.

We’ll keep it simple with the following mechanics/assumptions:
- We have a fixed number (adjustable parameter) of agents that are initially assigned random cells in the grid.
- A fixed percentage (adjustable parameter) of those agents are masked/unmasked throughout the entire simulation.
- At every step, an agent either stays at his/her cell or moves to an adjacent one.
- When a susceptible agent is in the cell with another infected agent, he/she may become infected based on a probability (adjustable parameter) that differs between masked and unmasked agents.
- If an agent is infected, he/she recovers and becomes immune after a certain number of simulation steps (adjustable parameter).
- Immunity goes away after a certain number of steps (adjustable parameter) and the agent becomes susceptible once again.
Visuals
Our agents will be represented as circles in our grid.
Susceptible agents are blue, infected are red, recovered & immune are green…
If a circle is (not) filled it means the agent is (not) masked…

Additionally, we’ll add a section for all adjustable parameters in our model.

Finally, we’ll have line charts dynamically showing us the number of susceptible, infected, and recovered agents throughout the simulation…

Code
In order to serve our agent-based model and the associated visualizations, we’ll be writing 2 Python scripts.
The first script which I’ve name covid_model.py has the Mesa model details that serves as the backend logic.
The 2nd script named covid_model_visualization.py will import the model parameters and the Agent and CovidModel classes from our first script. It will then serve our model and visualizations…
Final output
We now how our agent-based model ready.
We can adjust the parameters on the web page, run our Simulation, and visualize how are agents are interacting along with the counts of susceptible, infected, and recovered agents.

Note that you can click on "Step" to run a single step of the simulation or you can adjust the frames per second and click on "Start" to run the simulation at your desired speed.