def test_disease_progresses_with_contact_si(self):
num_steps = 10
population_size = 5
# Set up a population that is well-connected (here, totally connected).
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
graph.add_edges_from(complete_graph_edge_list(population_size))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=1.0,
num_treatments=0,
max_treatments=10,
initial_health_state=[
0 if i % 2 == 0 else 1 for i in range(graph.number_of_nodes())
],
)
agent = random_agents.RandomAgent(env.action_space, lambda x: 0,
env.observation_space)
initial_state = copy.deepcopy(env.state)
# Ensure that there are more infected people after running the simulation
# for some time.
test_util.run_test_simulation(env=env,
agent=agent,
num_steps=num_steps)
self.assertGreater(
num_in_health_state(env.state,
env.state_name_to_index["infected"]),
num_in_health_state(initial_state,
env.state_name_to_index["infected"]),
)
def test_last_population_used_for_infection_propagation(self):
population_size = 3
seed = 1
# Set up a graph I - S - S with a 1.0 infection rate. If we strictly use
# the 'frozen' population state to propagate disease, only the middle node
# should become infected, which is the desired behavior. If we use the
# evolving population and iterate from left to right, we should also expect
# the rightmost node to become infected, which is not the desired behavior.
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
graph.add_edges_from([(0, 1), (1, 2)])
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=1.0,
num_treatments=0,
max_treatments=population_size,
initial_health_state=[1, 0, 0],
)
env.set_scalar_reward(lambda x: 0)
env.seed(seed)
# The rightmost person should be SUSCEPTIBLE at t=0.
self.assertEqual(env.state_name_to_index["susceptible"],
env.state.health_states[2])
# The rightmost person should remain SUSCEPTIBLE at t=1.
env.step(np.arange(population_size)) # Treatment has no effect.
self.assertEqual(env.state_name_to_index["susceptible"],
env.state.health_states[2])
def test_disease_does_not_progress_without_contact(self):
num_steps = 10
# Set up an environment with some infected people but no contact between
# members of the population.
graph = nx.Graph()
graph.add_nodes_from(range(50))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.5,
num_treatments=0,
max_treatments=10,
initial_health_state=[0 for _ in graph],
)
agent = random_agents.RandomAgent(env.action_space, lambda x: 0,
env.observation_space)
initial_health_state = get_population_health_state(env.state)
# Run the simulation and ensure that the population's health state does
# change (because there's no opportunity for disease to spread due to
# the abscence of contact between people).
test_util.run_test_simulation(env=env,
agent=agent,
num_steps=num_steps)
final_health_state = get_population_health_state(env.state)
self.assertEqual(initial_health_state, final_health_state)
def test_disease_progresses_with_contact_during_burn_in(self):
population_size = 10
burn_in = 5
# Construct a chain graph.
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
for i in range(1, population_size):
graph.add_edge(i - 1, i)
# Construct a SI environment where the left-most node is infected.
initial_health_state = [1] + [0 for _ in range(population_size - 1)]
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=1.0,
num_treatments=0,
max_treatments=population_size,
initial_health_state=initial_health_state,
burn_in=burn_in,
)
# Before burn in, the initial health state should remain unchanged.
self.assertCountEqual([1] + [0 for _ in range(population_size - 1)],
env.state.health_states)
# After burn-in, the infection frontier should have moved by burn_in steps.
env.reset()
self.assertEqual(
[1 for _ in range(burn_in + 1)] +
[0 for _ in range(population_size - burn_in - 1)],
env.state.health_states,
)
def test_disease_prevalence_correct(self):
num_steps = 4
population_size = 40
healthy_state = 0
infectious_state = 1
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.0,
num_treatments=0,
max_treatments=population_size,
initial_health_state=[
healthy_state if i % 2 == 0 else infectious_state
for i in range(graph.number_of_nodes())
],
)
# Infection rates shouldn't change, so the most recent infection rate should
# be the same at each step.
metric = infectious_disease_metrics.DiseasePrevalence(env)
for _ in range(num_steps):
env.step(np.arange(population_size))
self.assertEqual(0.5, metric.measure(env))
def instantiate_environment_and_agent(
agent_class,
population_graph,
initial_health_state,
infection_probability=0.5,
num_treatments=5,
max_treatments=10,
seed=100,
agent_seed=50,
):
env = infectious_disease.build_si_model(
population_graph=population_graph,
infection_probability=infection_probability,
num_treatments=num_treatments,
initial_health_state=initial_health_state,
max_treatments=max_treatments,
)
agent = agent_class(
env.action_space,
rewards.NullReward(),
env.observation_space,
infectious_disease_agents.env_to_agent_params(env.initial_params),
)
env.seed(seed)
agent.seed(agent_seed)
_ = env.reset()
return env, agent
def test_render_does_not_raise(self):
graph = nx.Graph()
graph.add_nodes_from(range(50))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.5,
num_treatments=0,
max_treatments=10,
initial_health_state=[0 for _ in graph],
)
env.render({0: "blue", 1: "red"})
def test_render_fails_with_invalid_mode(self):
graph = nx.Graph()
graph.add_nodes_from(range(50))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.5,
num_treatments=0,
max_treatments=10,
initial_health_state=["susceptible" for _ in graph],
)
with self.assertRaises(ValueError):
env.render({0: "blue", 1: "red"}, mode="an invalid mode")
def test_rng_initialized_correctly(self):
# Set up an environment with some infected people but no contact between
# members of the population.
graph = nx.Graph()
graph.add_nodes_from(range(50))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.5,
num_treatments=0,
max_treatments=10,
initial_health_state=[0 for _ in graph],
)
env.state.rng.rand(
) # This will fail if rng is not initialized correctly.
def test_graph_space_contains_contact_network(self):
population_size = 50
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
graph.add_edges_from(complete_graph_edge_list(population_size))
env = infectious_disease.build_si_model(population_graph=graph,
infection_probability=1.0,
num_treatments=0,
max_treatments=10)
space = graph_space.GraphSpace(population_size, directed=False)
self.assertTrue(space.contains(env.state.population_graph))
def test_disease_permanently_eradicated_if_whole_population_treated(self):
population_size = 50
seed = 1
infected_indices = {4, 6, 30, 44}
initial_health_state = [
1 if i in infected_indices else 0 for i in range(population_size)
]
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
graph.add_edges_from(complete_graph_edge_list(population_size))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=1.0,
num_treatments=population_size,
max_treatments=population_size,
initial_health_state=initial_health_state,
)
env.set_scalar_reward(lambda x: 0)
env.seed(seed)
self.assertGreater(
num_in_health_state(env.state,
env.state_name_to_index["infected"]), 0)
# Treat everyone and ensure the disease is gone.
total_treatment = np.arange(population_size)
env.step(total_treatment)
self.assertEqual(
0,
num_in_health_state(env.state,
env.state_name_to_index["infected"]))
# Once the disease has been eradicated, the population should remain
# disease-free.
env.state.num_treatments = 0 # Stop treatment.
for _ in range(50):
env.step(np.arange(population_size))
self.assertEqual(
0,
num_in_health_state(env.state,
env.state_name_to_index["infected"]))
def test_healthy_population_counted_correctly(self):
num_steps = 4
population_size = 25
healthy_state = 0
graph = nx.Graph()
graph.add_nodes_from(range(population_size))
env = infectious_disease.build_si_model(
population_graph=graph,
infection_probability=0.0,
num_treatments=0,
max_treatments=population_size,
initial_health_state=[healthy_state for _ in graph],
)
for _ in range(num_steps):
env.step(np.arange(population_size))
metric = infectious_disease_metrics.PersonStepsInHealthState(
env, healthy_state)
self.assertEqual(metric.measure(env), num_steps * population_size)
