Agent-based SIR model Python example

In this article I implement a SIR model in Python, using a library that I developed for agent-based modeling in Python. In other words, I use an existing framework to implement an agent-based SIR model.

Introducing agent-based SIR model

SIR stands for susceptible, infected, and recovered. As illustrated in below figure, in a SIR model, agents can experience the following state transitions: From susceptible to infected, from infected to recovered, and from recovered to susceptible.

general sir model

The state transitions illustrated in above figure are usually implemented with a random distribution and some specified probability. Furthermore, you do no have to allow for all state transitions. In this article I, for example, will not allow agents to fall back from their recovered state to a susceptible state. Below is the SIR model implemented in this article.

sir model used for this article´s agent-based sir model

In the same way you could also formulate a SIR model in which you allow for fall back from recovered state to infected state, not just the susceptible state.

Agent-based modeling framework in Python

I created a framework for agent-based modeling. The framework is meant to be expanded and upgraded over time. It currently supports a setup that allows agents, simulation environment and populations to be placed in a grid-based model. Simulation results can be written into a SQLite database in standardized format, and visualizations and animations are generated based on the data stored in that database. The framework comes in the form of several Python files (modules) that can be editted to add custom functionality.

framework for agent-based sir model

If you want to check out the framework you can follow it on GitHub. Here is the link:

The special thing to note here is that the framework supports grid-based agent-based modeling. Agents interact in a 2D grid. This 2D grid facilitates the simulation environment. One or several agents can be located in the same cell, depending on model settings. Agents can interact with other agents within a specified neighbourhood. In this article ineractions refer to disease transfer, resulting in disease spread.

grid for this article´s agent-based sir model

Agents have attributes and functionality, represented by inner dynamics, variables and states. The inner workings of an agent affects its actions, decisions and interactions with other agents. Those interactions can be influenced by the interaction environment, too.

I will now proceed by documenting relevant ABM model implementation in Python.

Implementing the agent-based SIR model in Python

In below code I make use of the agent-based modeling framework. I set up a simulation run that lasts for 300 iterations. The probability of susceptible agents being infected is 7%, for every encounter with an infected agent. For every iteration an infected agent, on the other hand, has a 3% probability of recovering.

__author__ = "Linnart Felkl"
__email__ = "LinnartSF@gmail.com"

if __name__ == "__main__":

    print("demo starts")

    import sys
    from pathlib import Path
    file = Path(__file__).resolve()
    parent, root = file.parent, file.parents[1]
    sys.path.append(str(root))

    # remove the current file's directory from sys.path, unless already removed
    try:
        sys.path.remove(str(parent))
    except ValueError:
        pass

    import data
    import stats
    import config
    import framework
    import random
    import animation

    # setup database manager and connection
    db = data.Database("sqlite3", config.path_databasefile)
    db_manager = data.Manager(db)

    # create an empty environment
    env = framework.Environment(1, True, 20, 20, db_manager) # Environment constructor implicitly resets environment table in database

    # create initial population of healthy humans
    pops = framework.Populations(amount = 1, env = env, db_manager = db_manager, attributes = ["infected","recovered"], datatypes = ["INTEGER","INTEGER"])
    pops.add_population(name = "humans", 
                        size = 200, 
                        attributes = ["infected","recovered"], 
                        datatypes = ["INTEGER","INTEGER"], 
                        initialvals = [0, 0]
                        )

    # randomly infect 5%
    pop = pops.Populations["humans"]
    sampleagents = pop.get_agents(int(0.05*pop.Size))
    for agent in sampleagents: agent.Attributes["infected"] = 1

    _prob_infection = 0.07
    _prob_recovery = 0.03
    
    # setup simulation
    sim = framework.Simulation(300)

    # make sure that environment and agents tables in database are setup at this time
    pops.write_env_to_db(sim.Iteration)
    pops.write_agents_to_db(sim.Iteration)

    # execute simulation run, store every 10th run into the database
    while sim.run():

        # look at every agent
        for agent in pop.get_agents():
            if agent.get_attr_value("infected") == 1:
                
                neighbours = env.get_neighbourhood(agent)
                for neighbour in neighbours:
                    if neighbour.get_attr_value("infected") == 0 and neighbour.get_attr_value("recovered") == 0:
                        
                        # this neighbour is not resistant and not infected; infect with specified probability
                        if random.uniform(0, 1) < _prob_infection: neighbour.set_attr_value("infected", 1)
            
                # the infected agent recovers with a specified probability
                if random.uniform(0, 1) < _prob_recovery: 
                    agent.set_attr_value("recovered", 1)
                    agent.set_attr_value("infected", 0)
                
        # update results in database, for agents and for environment
        pops.write_agents_to_db(sim.Iteration)
        pops.write_env_to_db(sim.Iteration)
        pops.write_density_to_db(sim.Iteration)

    # get dataframes with simulation results 
    humans_df = db_manager.get_populationdf(pop = "humans")
    env_df = db_manager.get_environmentdf()
    density_df = db_manager.get_densitydf()
    
    # visualize simulation data
    stats.set_fontsizes(8,10,12)

    stats.plot_agentattr_lines("infected", humans_df)
    stats.save_plot("infection_curves")
    
    stats.plot_agentattr_lines("recovered", humans_df)
    stats.save_plot("recovery_curves")

    stats.plot_grid_occupation(env_df, ["humans"])
    stats.save_plot("human_locations")

    stats.plot_density_markersize(density_df, "infected", 50, "red")
    stats.save_plot("infection_density")

    stats.plot_density_markersize(density_df, "recovered", 50, "green")
    stats.save_plot("recovery_density")

    stats.plot_avgattr_lines(["infected","recovered"], humans_df)
    stats.save_plot("avginfectedavgrecovered")

    # create and save animations
    animation.animate_density(
        df = density_df,
        filename = "infectionanimation",
        attr = "infected",
        defaultsize = 50,
        color = "red",
        tpf = 0.05
    )

    animation.animate_density(
        df = density_df,
        filename = "recoveryanimation",
        attr = "recovered",
        defaultsize = 50,
        color = "green",
        tpf = 0.05
    )

    # end program
    db.close()
    print("demo ends")

Above agent-based SIR model implementation generates results that are stored in a SQLite database. This database is then lateron used for visualizations. These visualizations can be implemented in a separate script, by another analyst, and at a different point in time. I.e. model execution and data visualization are two separate activities that are decoupled by the database.

Agent-based SIR model simulation results

I can now finally review the results that were generated by the grid-based agent-based SIR model. Below is a visualization of agent locations throughout this simulation run.

Below animation shows which agents were infected throughout time. The animation furthermore shows agent locations. In this simple example the agent locations are static.

Next, an animation of agent recovery throughout simulation time.

Below chart visualizes the mean likelihood of an agent being infected or recovered, at the given simulation time.

Agent recovery ensures that not all agents are infected. However, agents in this model have static locations. In another simulation run we could allow for dynamic agent locations, or even some sort of diffusion / migration. This would most likely increase the total amount of agents that are infected throughout the simulation run.

The agent-based modeling framework applied in this article can deliver additional charts. Moreover, the simulation results remain stored in a SQLite database that could be analyzed in any custom way. This would however require some additional customization/adaption, i.e. coding.

Final remarks and related content

Agent-based simulation is a simulation method that seeks to explain of complex macroscopic system behaviour by describing agent dynamics and interactions on the microscopic system level. It is a method that is rather abstract and is often used to understand complex systems. Complex systems are, well, more complex than complicated systems, and, on the other hand, less complex than chaotic systems. The US economcy, visible on the macroscopic level in the form of e.g. gross-domestic product, is one example of a complex system. Agent-based models can e.g. be used to understand the emergence of certain phenomena and trends in US economcy.

In this article I implemented a recently very popular model for disease spread. This is in fact another interesting area in which agent-based modeling is frequently applied. Using a framework for agent-based modeling I implemented an agent-based SIR model. The framework’s visualization and animation capabilities were used to review simulation results. The framework itself can be used to implement other types of agent-based simulation models – and the framework will be upgraded and regularly expanded over time.

Here are some other SCDA related to agent-based simulation:

You May Also Like

Leave a Reply

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.