Note

You can download this demonstration as a Jupyter Notebook here

Virus spread

This notebook presents an agent-based model that simulates the propagation of a disease through a network. It demonstrates how to use the agentpy package to create and visualize networks, use the interactive module, and perform different types of sensitivity analysis.

# Model design
import agentpy as ap
import networkx as nx
import random

# Visualization
import matplotlib.pyplot as plt
import seaborn as sns
import IPython

About the model

The agents of this model are people, which can be in one of the following three conditions: susceptible to the disease (S), infected (I), or recovered (R). The agents are connected to each other through a small-world network of peers. At every time-step, infected agents can infect their peers or recover from the disease based on random chance.

Defining the model

We define a new agent type Person by creating a subclass of Agent. This agent has two methods: setup() will be called automatically at the agent’s creation, and being_sick() will be called by the Model.step() function. Three tools are used within this class:

• Agent.p returns the parameters of the model

• Agent.neighbors() returns a list of the agents’ peers in the network

• random.random() returns a uniform random draw between 0 and 1

class Person(ap.Agent):

    def setup(self):
        """ Initialize a new variable at agent creation. """
        self.condition = 0  # Susceptible = 0, Infected = 1, Recovered = 2

    def being_sick(self):
        """ Spread disease to peers in the network. """
        for n in self.neighbors():
            if n.condition == 0 and self.p.infection_chance > random.random():
                n.condition = 1  # Infect susceptible peer
        if self.p.recovery_chance > random.random():
            self.condition = 2  # Recover from infection

Next, we define our model VirusModel by creating a subclass of Model. The four methods will be called automatically, as described in Running a simulation.

class VirusModel(ap.Model):

    def setup(self):
        """ Initializes the agents and network of the model. """

        self.p.population = p = int(self.p.population)
        # Prepare a small-world network graph
        graph = nx.watts_strogatz_graph(p,
                                        self.p.number_of_neighbors,
                                        self.p.network_randomness)

        # Create agents and network
        self.add_agents(p, Person)
        self.add_network(graph=graph, agents=self.agents)

        # Infect a random share of the population
        I0 = int(self.p.initial_infections * self.p.population)
        self.agents.random(I0).condition = 1

    def update(self):
        """ Records variables after setup and each step. """
        # Record share of agents with each condition
        for i, c in enumerate(('S', 'I', 'R')):
            self[c] = (len(self.agents.select(self.agents.condition == i))
                       / self.p.population)
            self.record(c)

        # Stop simulation if disease is gone
        if self.I == 0:
            self.stop()

    def step(self):
        """ Defines the models' events per simulation step. """
        # Call 'being_sick' for infected agents
        self.agents(self.agents.condition==1).being_sick()

    def end(self):
        """ Records evaluation measures at the end of the simulation. """
        # Record final evaluation measures
        self.measure('Total share infected', self.I + self.R)
        self.measure('Peak share infected', max(self.log['I']))

Running a simulation

To run our model, we define a dictionary with our parameters. We then create a new instance of our model, passing the parameters as an argument, and use the method Model.run() to perform the simulation and return it’s output.

parameters = {
    'population': 1000,
    'infection_chance': 0.3,
    'recovery_chance': 0.1,
    'initial_infections': 0.1,
    'number_of_neighbors': 2,
    'network_randomness': 0.5
}

model = VirusModel(parameters)
results = model.run()

Completed: 75 steps
Run time: 0:00:00.420014
Simulation finished

Analyzing results

The simulation returns a DataDict of recorded data with dataframes:

results

DataDict {
'log': Dictionary with 4 keys
'parameters': Dictionary with 6 keys
'measures': DataFrame with 2 variables and 1 row
'variables': DataFrame with 3 variables and 76 rows
}

To visualize the evolution of our variables over time, we create a plot function.

def virus_stackplot(data, ax):
    """ Stackplot of people's condition over time. """
    x = data.index.get_level_values('t')
    y = [data[var] for var in ['I', 'S', 'R']]

    sns.set()
    ax.stackplot(x, y, labels=['Infected', 'Susceptible', 'Recovered'],
                 colors = ['r', 'b', 'g'])

    ax.legend()
    ax.grid(False)
    ax.set_xlim(0, max(1, len(x)-1))
    ax.set_ylim(0, 1)
    ax.set_xlabel("Time steps")
    ax.set_ylabel("Percentage of population")

fig, ax = plt.subplots()
virus_stackplot(results.variables, ax)

Creating an animation

We can also animate the model’s dynamics as follows. The function animation_plot() takes a model instance and displays the previous stackplot together with a network graph. The function animate() will call this plot function for every time-step and return an matplotlib.animation.Animation.

def animation_plot(m, axs):
    ax1, ax2 = axs
    ax1.set_title("Virus spread")
    ax2.set_title(f"Share infected: {m.I}")

    # Plot stackplot on first axis
    virus_stackplot(m.output.variables, ax1)

    # Plot network on second axis
    color_dict = {0:'b', 1:'r', 2:'g'}
    colors = [color_dict[c] for c in m.agents.condition]
    nx.draw_circular(m.env.graph, node_color=colors,
                     node_size=50, ax=ax2)

fig, axs = plt.subplots(1, 2, figsize=(8, 4)) # Prepare figure
parameters['population'] = 50 # Lower population for better visibility
animation = ap.animate(VirusModel(parameters), fig, axs, animation_plot)

Using Jupyter, we can display this animation directly in our notebook.

IPython.display.HTML(animation.to_jshtml())

Once Loop Reflect