def test_resistances_not_same_element():
"""Tests that when we generate the node resistance we get a different quantity for each node instead of the same value for each node.
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test, initial_infected=1, pre_gen_data=100)
assert my_data.epi_data[(0, 0)]["Resistance"] != my_data.epi_data[(
0, 1)]["Resistance"]
def test_update_infection_stage_tuple():
"""We test that when we update an infection stage that:
1) The update goes through correctly
2) Only the intended node is getting updated
3) The status change is correctly added to the log """
my_data = epidemic_data(G_lattice,
initial_infected=[(0, 0)],
pre_gen_data=100)
my_data.update_infection_stage([(1, 1), (2, 2)], "Stage 1", 1)
my_data.update_infection_stage([(2, 2)], "Stage 2", 2)
assert my_data.epi_data[(1, 1)]["Infection Stage"] == "Stage 1"
assert my_data.epi_data[(1, 1)]["Infection Stage Started"] == 1
assert my_data.epi_data[(2, 2)]["Infection Stage"] == "Stage 2"
assert my_data.epi_data[(2, 2)]["Infection Stage Started"] == 2
assert my_data.epi_data[(1, 1)]["History"]["Infection Stage Log"] == [
"Susceptible", "Stage 1"
]
assert my_data.epi_data[(2, 2)]["History"]["Infection Stage Log"] == [
"Susceptible", "Stage 1", "Stage 2"
]
assert my_data.epi_data[(1,
1)]["History"]["Infection Stage Times"] == [0, 1]
assert my_data.epi_data[(2, 2)]["History"]["Infection Stage Times"] == [
0, 1, 2
]
def test_initialise_data_structure():
"""Tests that the output is as expected for two types of networks.
"""
G_test = nx.complete_graph(10)
my_data = epidemic_data(G_test, initial_infected=[0], pre_gen_data=100)
assert 0 in my_data.epi_data
assert my_data.epi_data[1]["Infection Stage"] == "Susceptible"
def test_update_exposure_level_tuple():
"""We call the method which increments a nodes exposure and check that the correct entry in the dictionary is updated and only the correct entry"""
my_data_structure = epidemic_data(G_lattice,
initial_infected=[(1, 1)],
pre_gen_data=100)
my_data_structure.update_exposure_level((2, 2), 3)
assert my_data_structure.epi_data[(2, 2)]["Exposure Level"] == 3
assert my_data_structure.epi_data[(3, 3)]["Exposure Level"] == 0
def test_initialise_data_structure_infected_as_int():
"""Tests that the output is as expected for two types of networks.
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test, initial_infected=4, pre_gen_data=100)
assert my_data.epi_data[(1, 1)]["Infection Stage"] == "Infected"
assert my_data.epi_data[(0, 0)]["Infection Stage"] == "Infected"
def test_update_exposure_level():
"""We call the method which increments a nodes exposure and check that the correct entry in the dictionary is updated"""
my_data_structure = epidemic_data(G_complete,
initial_infected=[1],
pre_gen_data=100)
my_data_structure.update_exposure_level(2, 3)
assert my_data_structure.epi_data[2]["Exposure Level"] == 3
assert my_data_structure.epi_data[1]["Exposure Level"] == 0
assert my_data_structure.epi_data[3]["Exposure Level"] == 0
def test_initialise_data_structure_tuple_nodes():
"""Tests that the output is as expected for two types of networks.
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test,
initial_infected=[(0, 0)],
pre_gen_data=100)
assert (1, 1) in my_data.epi_data
assert my_data.epi_data[(1, 1)]["Infection Stage"] == "Susceptible"
def test_2_par_distribution_tuple():
def dist_2_par(par_1, par_2, n):
return np.array([5] * n) * par_1 * par_2
my_data = epidemic_data(G_lattice,
initial_infected=[(1, 1)],
pre_gen_data=100,
infection_period_distribution=dist_2_par,
infection_period_parameters=[1, 2])
assert my_data.epi_data[(1, 1)]["Infection Period"] == 10
def test_update_infection_stage():
"""We test that the dictionary updater updates current values correctly, we are not checking the history yet
"""
G_test = nx.complete_graph(10)
my_data = epidemic_data(G_test, initial_infected=[1], pre_gen_data=100)
my_data.update_infection_stage([0], "Infected", 10)
assert my_data.epi_data[0]["Infection Stage"] == "Infected"
assert my_data.epi_data[0]["Infection Stage Started"] == 10
assert my_data.epi_data[2]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[2]["Infection Stage Started"] == 0
def test_3_par_distribution():
def dist_3_par(par_1, par_2, par_3, n):
return np.array([5] * n) * par_1 * par_2 * par_3
my_data = epidemic_data(G_complete,
initial_infected=[1],
pre_gen_data=100,
infection_period_distribution=dist_3_par,
infection_period_parameters=[1, 2, 3])
assert my_data.epi_data[1]["Infection Period"] == 30
def test_initialise_resistances():
"""Cant think of a better way to check that the resistances are exponential(1) w/out resorting to monte carlo.
We also check that it's not the same in memory.
"""
my_data = epidemic_data(G_complete, initial_infected=[0], pre_gen_data=100)
assert all([
my_data.epi_data[node]["Resistance"] > 0 for node in my_data.epi_data
if my_data.epi_data[node]["Infection Stage"] == "Susceptible"
])
assert my_data.epi_data[1]["Resistance"] != my_data.epi_data[2][
"Resistance"]
def test_update_infection_stage_list_input():
"""We test that the dictionary updater updates the correct values
and only the correct values.
"""
G_test = nx.complete_graph(10)
my_data = epidemic_data(G_test, initial_infected=[0], pre_gen_data=100)
my_data.update_infection_stage([1, 2], "Infected", 10)
assert my_data.epi_data[1]["Infection Stage"] == "Infected"
assert my_data.epi_data[1]["Infection Stage Started"] == 10
assert my_data.epi_data[2]["Infection Stage"] == "Infected"
assert my_data.epi_data[2]["Infection Stage Started"] == 10
def test_update_infection_stage_tuple_nodes():
"""We test that the dictionary updater updates current values correctly, we are not checking the history yet
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test,
initial_infected=[(1, 1)],
pre_gen_data=100)
my_data.update_infection_stage([(0, 0)], "Infected", 10)
assert my_data.epi_data[(0, 0)]["Infection Stage"] == "Infected"
assert my_data.epi_data[(0, 0)]["Infection Stage Started"] == 10
assert my_data.epi_data[(0, 1)]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[(0, 1)]["Infection Stage Started"] == 0
def test_update_infection_stage_list_input_tuple_nodes():
"""We test that the dictionary updater updates the correct values
and only the correct values.
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test,
initial_infected=[(1, 1)],
pre_gen_data=100)
my_data.update_infection_stage([(0, 0), (0, 1)], "Infected", 10)
assert my_data.epi_data[(0, 0)]["Infection Stage"] == "Infected"
assert my_data.epi_data[(0, 0)]["Infection Stage Started"] == 10
assert my_data.epi_data[(0, 1)]["Infection Stage"] == "Infected"
assert my_data.epi_data[(0, 1)]["Infection Stage Started"] == 10
def __init__(self,
G,
beta,
infection_period_parameters,
initial_infected,
time_increment,
max_iterations,
hazard_rate=None,
infection_period_distribution=None,
SIS=False,
increment_network=None,
custom_behaviour=None):
"""This class manages the simulation of the epidemic and the simulation of the dynamic network (if the network is dynamic).
If the network is static, then
Arguments:
G {Networkx.graph} -- A NetworkX graph object
beta {float} -- The thinning parameter
infection_period_parameters {list} -- A list of parameters for the infection period distribution
initial_infected {int, list} -- [description]
time_increment {float} -- The length of the time step of the simulation
max_iterations {int} -- The maximum number of iterations that will be performed
Keyword Arguments:
hazard_rate {function} -- A function of the form f(x) (default: f(x) = 1)
infection_period_distribution {function} -- A numpy random number distribution (default: {None})
SIS {bool} -- Boolean on whether the epidemic is SIS, if not it will be treated as SIR (default: False)
increment_network {method} -- A method of the form increment_network(increment_length). This method will be called during the simulation to move the network forward by the network_increment.
custom_behaviour {function} -- Allows users to execute custom behaviour during the simulation. This is useful for customising the simulation to your own purposes, such as treatment scenarios. (default: {None})
TODO: Remove the beta parameter, too confusing
"""
self.G = G
self.beta = beta
self.infection_period_parameters = infection_period_parameters
self.inf_starting = initial_infected
self.time_increment = time_increment
self.max_iterations = max_iterations
self.hazard_rate = hazard_rate
self.SIS = SIS
self.increment_network = increment_network
self.custom_behaviour = custom_behaviour
self.time = 0
self.N = nx.number_of_nodes(self.G)
self.data_structure = epidemic_data(G, initial_infected, 100,
infection_period_distribution,
infection_period_parameters)
self.epi_data = self.data_structure.epi_data
self.hazard = hazard_class(self.hazard_rate)
def test_update_infection_stage_2_updates():
"""We test that the dictionary updater updates the correct values
and only the correct values.
"""
G_test = nx.complete_graph(10)
my_data = epidemic_data(G_test, initial_infected=[1], pre_gen_data=100)
my_data.update_infection_stage([0], "Infected", 10)
my_data.update_infection_stage([0], "Susceptible", 20)
assert my_data.epi_data[0]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[0]["History"]["Infection Stage Times"] == [
0, 10, 20
]
assert my_data.epi_data[2]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[2]["History"]["Infection Stage Times"] == [0]
def test_update_infection_stage_2_updates_tuple_nodes():
"""We test that the dictionary updater updates the correct values
and only the correct values.
"""
G_test = nx.grid_2d_graph(2, 2)
my_data = epidemic_data(G_test,
initial_infected=[(1, 1)],
pre_gen_data=100)
my_data.update_infection_stage([(0, 0)], "Infected", 10)
my_data.update_infection_stage([(0, 0)], "Susceptible", 20)
assert my_data.epi_data[(0, 0)]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[(0, 0)]["History"]["Infection Stage Times"] == [
0, 10, 20
]
assert my_data.epi_data[(0, 1)]["Infection Stage"] == "Susceptible"
assert my_data.epi_data[(0, 1)]["History"]["Infection Stage Times"] == [0]
def test_data_pre_gen():
def dist_1_par(par_1, n):
return np.array([5] * n) * par_1
my_data = epidemic_data(G_lattice,
pre_gen_data=100,
initial_infected=[(1, 1)],
infection_period_distribution=dist_1_par,
infection_period_parameters=1)
test_resistance = my_data.epi_data[(1,
1)]["Pre-generated Data"]["Resistance"]
test_infection_periods = my_data.epi_data[(
1, 1)]["Pre-generated Data"]["Infection Period"]
assert len(test_resistance) == 100
assert len(test_infection_periods) == 100
assert all([test_infection_periods[i] == 5 for i in range(100)]) == True
