Note

You can download this demonstration as a Jupyter Notebook here

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.

[1]:

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

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


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:

[2]:

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.

[3]:

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

# 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.

[4]:

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:

[5]:

results

[5]:

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.

[6]:

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.

[7]:

def animation_plot(m, axs):
ax1, ax2 = axs
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.

[8]:

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

[8]: