class TestSingleNetworkGrid(unittest.TestCase):
GRAPH_SIZE = 10
def setUp(self):
'''
Create a test network grid and populate with Mock Agents.
'''
G = nx.complete_graph(TestSingleNetworkGrid.GRAPH_SIZE)
self.space = NetworkGrid(G)
self.agents = []
for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
a = MockAgent(i, None)
self.agents.append(a)
self.space.place_agent(a, pos)
def test_agent_positions(self):
'''
Ensure that the agents are all placed properly.
'''
for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
a = self.agents[i]
assert a.pos == pos
def test_get_neighbors(self):
assert len(self.space.get_neighbors(
0, include_center=True)) == TestSingleNetworkGrid.GRAPH_SIZE
assert len(self.space.get_neighbors(
0, include_center=False)) == TestSingleNetworkGrid.GRAPH_SIZE - 1
def test_move_agent(self):
initial_pos = 1
agent_number = 1
final_pos = TestSingleNetworkGrid.GRAPH_SIZE - 1
_agent = self.agents[agent_number]
assert _agent.pos == initial_pos
assert _agent in self.space.G.node[initial_pos]['agent']
assert _agent not in self.space.G.node[final_pos]['agent']
self.space.move_agent(_agent, final_pos)
assert _agent.pos == final_pos
assert _agent not in self.space.G.node[initial_pos]['agent']
assert _agent in self.space.G.node[final_pos]['agent']
def test_is_cell_empty(self):
assert not self.space.is_cell_empty(0)
assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1)
def test_get_cell_list_contents(self):
assert self.space.get_cell_list_contents([0]) == [self.agents[0]]
assert self.space.get_cell_list_contents(
list(range(TestSingleNetworkGrid.GRAPH_SIZE))) == [
self.agents[0], self.agents[1], self.agents[2]
]
def test_get_all_cell_contents(self):
assert self.space.get_all_cell_contents() == [
self.agents[0], self.agents[1], self.agents[2]
]
class TestSingleNetworkGrid(unittest.TestCase):
GRAPH_SIZE = 10
def setUp(self):
'''
Create a test network grid and populate with Mock Agents.
'''
G = nx.complete_graph(TestSingleNetworkGrid.GRAPH_SIZE)
self.space = NetworkGrid(G)
self.agents = []
for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
a = MockAgent(i, None)
self.agents.append(a)
self.space.place_agent(a, pos)
def test_agent_positions(self):
'''
Ensure that the agents are all placed properly.
'''
for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
a = self.agents[i]
assert a.pos == pos
def test_get_neighbors(self):
assert len(self.space.get_neighbors(0, include_center=True)) == TestSingleNetworkGrid.GRAPH_SIZE
assert len(self.space.get_neighbors(0, include_center=False)) == TestSingleNetworkGrid.GRAPH_SIZE - 1
def test_move_agent(self):
initial_pos = 1
agent_number = 1
final_pos = TestSingleNetworkGrid.GRAPH_SIZE - 1
_agent = self.agents[agent_number]
assert _agent.pos == initial_pos
assert _agent in self.space.G.node[initial_pos]['agent']
assert _agent not in self.space.G.node[final_pos]['agent']
self.space.move_agent(_agent, final_pos)
assert _agent.pos == final_pos
assert _agent not in self.space.G.node[initial_pos]['agent']
assert _agent in self.space.G.node[final_pos]['agent']
def test_is_cell_empty(self):
assert not self.space.is_cell_empty(0)
assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1)
def test_get_cell_list_contents(self):
assert self.space.get_cell_list_contents([0]) == [self.agents[0]]
assert self.space.get_cell_list_contents(list(range(TestSingleNetworkGrid.GRAPH_SIZE))) == [self.agents[0],
self.agents[1],
self.agents[2]]
def test_get_all_cell_contents(self):
assert self.space.get_all_cell_contents() == [self.agents[0],
self.agents[1],
self.agents[2]]
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,
):
self.num_nodes = num_nodes
prob = avg_node_degree / self.num_nodes
self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
self.grid = NetworkGrid(self.G)
self.schedule = RandomActivation(self)
self.initial_outbreak_size = (initial_outbreak_size
if initial_outbreak_size <= num_nodes
else num_nodes)
self.virus_spread_chance = virus_spread_chance
self.virus_check_frequency = virus_check_frequency
self.reovery_chance = recovery_chance
self.gain_resistance_chance = gain_resistance_chance
self.datacollector = DataCollector({
'Infected': number_infected,
'Susceptible': number_susceptible,
'Resistant': number_resistant,
})
# Create agents
for i, node in enumerate(self.G.nodes()):
a = VirusAgent(
i,
self,
State.SUSCEPTIBLE,
self.virus_spread_chance,
self.virus_check_frequency,
self.reovery_chance,
self.gain_resistance_chance,
)
self.schedule.add(a)
# Add the agent to the node
self.grid.place_agent(a, node)
# Infect some agents
infected_nodes = self.random.sample(self.G.nodes(),
self.initial_outbreak_size)
for a in self.grid.get_all_cell_contents(infected_nodes):
a.state = State.INFECTED
self.running = True
self.datacollector.collect(self)
def resistant_succeptible_ratio(self):
try:
return number_state(self, State.RESISTANT) / number_state(
self, State.SUSCEPTIBLE)
except ZeroDivisionError:
return math.inf
def step(self):
self.schedule.step()
# collect data
self.datacollector.collect(self)
def run_model(self, n):
for i in range(n):
self.step()
class Population(Model):
"""Population
Adapted from https://www.medrxiv.org/content/10.1101/2020.03.18.20037994v1.full.pdf
Model Parameters:
spread_chance: probability of infection based on contact
gamma: mean incubation period
alpha: probability of become asymptomatic vs symptomatic
gamma_AR: infectious period for asymptomatic people
gamma_YR: infectious period for symptomatic people
delta: death rate due to disease
The social distancing func takes time passed -> new interaction multiplier
"""
def __init__(self,
graph,
model_parameters,
social_distancing_func=lambda x: 1):
# Model initialization
self.population_size = model_parameters['population size']
self.initial_outbreak_size = model_parameters['initial outbreak size']
self.graph = graph
self.grid = NetworkGrid(self.graph)
self.schedule = SimultaneousActivation(self)
self.social_distancing_func = social_distancing_func
self.datacollector = DataCollector({ # "Exposed": count_exposed,
"Susceptible": count_susceptible,
"Recovered": count_recovered,
# "Asymptomatic": count_asymptomatic,
# "Symptomatic": count_symptomatic,
"Diseased": count_diseased,
"Removed": count_removed
})
self.model_parameters = model_parameters
for i, node in enumerate(self.graph.nodes()):
a = Person(i, self, State.SUSCEPTIBLE, model_parameters,
social_distancing_func)
self.schedule.add(a)
self.grid.place_agent(a, i)
if i % 100 == 0:
logger.info("Finished with agent " + str(i))
infected_nodes = self.random.sample(self.graph.nodes(),
self.initial_outbreak_size)
for a in self.grid.get_cell_list_contents(infected_nodes):
a.status = State.EXPOSED
self.datacollector.collect(self)
print("Model initialized...\n")
def step(self):
self.datacollector.collect(self)
self.schedule.step()
def run(self, n):
for i in range(n):
logger.info("Steps Completed: " + str(i))
self.step()
def run_until_stable(self):
max_steps = 1e3
steps = 0
window_size = 50
while steps < max_steps:
self.step()
logger.info("Steps Completed:" + str(steps))
if steps > window_size:
data = self.get_data()
last_value = int(data.tail(1)['Diseased'])
if last_value == 0:
break
window_average = np.mean(
data.tail(window_size)
['Diseased']) # Window for determining stopping rule
if abs(last_value - window_average) / window_average < 0.005:
break
steps += 1
def cross_influence(self, influence_coefficients, ratios):
for inf_coeff, ratio in zip(influence_coefficients, ratios):
susceptibles = list(
filter(lambda x: x.status == State.SUSCEPTIBLE,
self.grid.get_all_cell_contents()))
to_flip = self.random.sample(
susceptibles, int(inf_coeff * ratio * len(susceptibles)))
for agent in to_flip:
agent.status = State.EXPOSED
def clear_social_distancing_func(self):
"""
Clears the social distancing function of the model and all its agents for pickling
:return: None
"""
self.social_distancing_func = None
for agent in self.grid.get_all_cell_contents():
agent.social_distancing_func = None
def reinstate_social_distancing_func(self,
social_distancing_func=lambda x: 1):
"""
Re-adds the social distancing func to the model and all its agents
:param social_distancing_func: social distancing func to be re-added
:return: None
"""
self.social_distancing_func = social_distancing_func
for agent in self.grid.get_all_cell_contents():
agent.social_distancing_func = social_distancing_func
def save_model(self, filename):
"""
Save the model to a pickle and dill file
:param filename: filename (without extension) to save to
:return: None
"""
with open(filename + ".dil", 'wb') as f:
dill.dump(self.social_distancing_func, f)
self.clear_social_distancing_func()
with open(filename + ".pkl", 'wb') as f:
pickle.dump(self, f)
def get_data(self):
return self.datacollector.get_model_vars_dataframe()
'''
class VirusModel(Model):
def __init__(self, num_nodes, avg_node_degree, initial_outbreak_size, alpha, beta, gamma, delta, k, n):
self.num_nodes = num_nodes
self.avg_node_degree = avg_node_degree
self.G = nx.barabasi_albert_graph(n=self.num_nodes,m=avg_node_degree)
self.grid = NetworkGrid(self.G)
self.schedule = RandomActivation(self)
self.initial_outbreak_size = initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
self.alpha = alpha
self.beta = beta
self.gamma = gamma
self.delta = delta
self.k=k
self.n=n
# Create agents
for i, node in enumerate(self.G.nodes()):
a = VirusAgent(i, self, State.SUSCEPTIBLE, self.alpha, self.beta, self.gamma, self.delta, self.k, self.n)
self.schedule.add(a)
# Add the agent to the node
self.grid.place_agent(a, node)
# Infect some nodes
active_nodes = random.sample(self.G.nodes(), self.initial_outbreak_size)
for a in self.grid.get_cell_list_contents(active_nodes):
a.state = State.ACTIVE
self.datacollector = DataCollector(
model_reporters={
"Infected": number_active,
"Susceptible": number_susceptible,
"Carrier": number_inactive,
"Removed": number_removed,
"Active Clustering": infective_clustering,
"Exposed Clustering": exposed_clustering,
"Infective Diffusion": infective_diffusion,
"Exposed Diffusion": exposed_diffusion
}
)
self.running = True
self.datacollector.collect(self)
def removed_susceptible_ratio(self):
try:
return number_state(self, State.REMOVED) / number_state(self, State.SUSCEPTIBLE)
except ZeroDivisionError:
return math.inf
def tyrant_remove(self):
if number_active(self) > 0:
if random.random() < 0.002:
actives = [a for a in self.grid.get_all_cell_contents() if a.state is State.ACTIVE]
node_for_removal = random.sample(actives, 1)
for a in node_for_removal:
a.state = State.REMOVED
# for a in self.grid.get_cell_list_contents(active_nodes):
# a.state = State.ACTIVE
def step(self):
self.tyrant_remove()
self.schedule.step()
self.datacollector.collect(self)
def run_model(self, n):
for i in range(n):
self.step()