Virus on a Network#
Summary#
This model is based on the NetLogo model “Virus on Network”. It demonstrates the spread of a virus through a network and follows the SIR model, commonly seen in epidemiology.
The SIR model is one of the simplest compartmental models, and many models are derivatives of this basic form. The model consists of three compartments:
S: The number of susceptible individuals. When a susceptible and an infectious individual come into “infectious contact”, the susceptible individual contracts the disease and transitions to the infectious compartment. I: The number of infectious individuals. These are individuals who have been infected and are capable of infecting susceptible individuals. R for the number of removed (and immune) or deceased individuals. These are individuals who have been infected and have either recovered from the disease and entered the removed compartment, or died. It is assumed that the number of deaths is negligible with respect to the total population. This compartment may also be called “recovered” or “resistant”.
For more information about this model, read the NetLogo’s web page: http://ccl.northwestern.edu/netlogo/models/VirusonaNetwork.
JavaScript library used in this example to render the network: d3.js.
Installation#
To install the dependencies use pip and the requirements.txt in this directory. e.g.
$ pip install -r requirements.txt
How to Run#
To run the model interactively, in this directory, run the following command
$ solara run app.py
Files#
model.py: Contains the agent class, and the overall model class.agents.py: Contains the agent class.app.py: Contains the code for the interactive Solara visualization.
Further Reading#
Stonedahl, F. and Wilensky, U. (2008). NetLogo Virus on a Network model. Center for Connected Learning and Computer-Based Modeling, Northwestern University, Evanston, IL.
Wilensky, U. (1999). NetLogo Center for Connected Learning and Computer-Based Modeling, Northwestern University, Evanston, IL.
Agents#
from enum import Enum
from mesa.discrete_space import FixedAgent
class State(Enum):
SUSCEPTIBLE = 0
INFECTED = 1
RESISTANT = 2
class VirusAgent(FixedAgent):
"""Individual Agent definition and its properties/interaction methods."""
def __init__(
self,
model,
initial_state,
virus_spread_chance,
virus_check_frequency,
recovery_chance,
gain_resistance_chance,
cell,
):
super().__init__(model)
self.state = initial_state
self.virus_spread_chance = virus_spread_chance
self.virus_check_frequency = virus_check_frequency
self.recovery_chance = recovery_chance
self.gain_resistance_chance = gain_resistance_chance
self.cell = cell
def try_to_infect_neighbors(self):
for agent in self.cell.neighborhood.agents:
if (agent.state is State.SUSCEPTIBLE) and (
self.random.random() < self.virus_spread_chance
):
agent.state = State.INFECTED
def try_gain_resistance(self):
if self.random.random() < self.gain_resistance_chance:
self.state = State.RESISTANT
def try_remove_infection(self):
# Try to remove
if self.random.random() < self.recovery_chance:
# Success
self.state = State.SUSCEPTIBLE
self.try_gain_resistance()
else:
# Failed
self.state = State.INFECTED
def check_situation(self):
if (self.state is State.INFECTED) and (
self.random.random() < self.virus_check_frequency
):
self.try_remove_infection()
def step(self):
if self.state is State.INFECTED:
self.try_to_infect_neighbors()
self.check_situation()
Model#
import math
import networkx as nx
import mesa
from mesa import Model
from mesa.discrete_space import CellCollection, Network
from mesa.examples.basic.virus_on_network.agents import State, VirusAgent
def number_state(model, state):
return sum(1 for a in model.grid.all_cells.agents if a.state is state)
def number_infected(model):
return number_state(model, State.INFECTED)
def number_susceptible(model):
return number_state(model, State.SUSCEPTIBLE)
def number_resistant(model):
return number_state(model, State.RESISTANT)
class VirusOnNetwork(Model):
"""A virus model with some number of agents."""
def __init__(
self,
num_nodes=10,
avg_node_degree=3,
initial_outbreak_size=1,
virus_spread_chance=0.4,
virus_check_frequency=0.4,
recovery_chance=0.3,
gain_resistance_chance=0.5,
seed=None,
):
super().__init__(seed=seed)
prob = avg_node_degree / num_nodes
graph = nx.erdos_renyi_graph(n=num_nodes, p=prob)
self.grid = Network(graph, capacity=1, random=self.random)
self.initial_outbreak_size = (
initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
)
self.datacollector = mesa.DataCollector(
{
"Infected": number_infected,
"Susceptible": number_susceptible,
"Resistant": number_resistant,
"R over S": self.resistant_susceptible_ratio,
}
)
VirusAgent.create_agents(
self,
num_nodes,
State.SUSCEPTIBLE,
virus_spread_chance,
virus_check_frequency,
recovery_chance,
gain_resistance_chance,
list(self.grid.all_cells),
)
# Infect some nodes
infected_nodes = CellCollection(
self.random.sample(list(self.grid.all_cells), self.initial_outbreak_size),
random=self.random,
)
for a in infected_nodes.agents:
a.state = State.INFECTED
self.running = True
self.datacollector.collect(self)
def resistant_susceptible_ratio(self):
try:
return number_state(self, State.RESISTANT) / number_state(
self, State.SUSCEPTIBLE
)
except ZeroDivisionError:
return math.inf
def step(self):
self.agents.shuffle_do("step")
# collect data
self.datacollector.collect(self)
App#
import math
import solara
from mesa.examples.basic.virus_on_network.model import (
State,
VirusOnNetwork,
number_infected,
)
from mesa.visualization import (
Slider,
SolaraViz,
SpaceRenderer,
make_plot_component,
)
from mesa.visualization.components import AgentPortrayalStyle
def agent_portrayal(agent):
node_color_dict = {
State.INFECTED: "red",
State.SUSCEPTIBLE: "green",
State.RESISTANT: "gray",
}
return AgentPortrayalStyle(color=node_color_dict[agent.state], size=20)
def get_resistant_susceptible_ratio(model):
ratio = model.resistant_susceptible_ratio()
ratio_text = r"$\infty$" if ratio is math.inf else f"{ratio:.2f}"
infected_text = str(number_infected(model))
return solara.Markdown(
f"Resistant/Susceptible Ratio: {ratio_text}<br>Infected Remaining: {infected_text}"
)
model_params = {
"seed": {
"type": "InputText",
"value": 42,
"label": "Random Seed",
},
"num_nodes": Slider(
label="Number of agents",
value=10,
min=10,
max=100,
step=1,
),
"avg_node_degree": Slider(
label="Avg Node Degree",
value=3,
min=3,
max=8,
step=1,
),
"initial_outbreak_size": Slider(
label="Initial Outbreak Size",
value=1,
min=1,
max=10,
step=1,
),
"virus_spread_chance": Slider(
label="Virus Spread Chance",
value=0.4,
min=0.0,
max=1.0,
step=0.1,
),
"virus_check_frequency": Slider(
label="Virus Check Frequency",
value=0.4,
min=0.0,
max=1.0,
step=0.1,
),
"recovery_chance": Slider(
label="Recovery Chance",
value=0.3,
min=0.0,
max=1.0,
step=0.1,
),
"gain_resistance_chance": Slider(
label="Gain Resistance Chance",
value=0.5,
min=0.0,
max=1.0,
step=0.1,
),
}
def post_process_lineplot(chart):
chart = chart.properties(
width=400,
height=400,
).configure_legend(
strokeColor="black",
fillColor="#ECE9E9",
orient="right",
cornerRadius=5,
padding=10,
strokeWidth=1,
)
return chart
model1 = VirusOnNetwork()
renderer = SpaceRenderer(model1, backend="altair")
renderer.draw_structure(
node_kwargs={"color": "black", "filled": False, "strokeWidth": 5},
edge_kwargs={"strokeDash": [6, 1]},
) # Do this to draw the underlying network and customize it
renderer.draw_agents(agent_portrayal)
# Plot components can also be in altair and support post_process
StatePlot = make_plot_component(
{"Infected": "red", "Susceptible": "green", "Resistant": "gray"},
backend="altair",
post_process=post_process_lineplot,
)
page = SolaraViz(
model1,
renderer,
components=[
StatePlot,
get_resistant_susceptible_ratio,
],
model_params=model_params,
name="Virus Model",
)
page # noqa