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_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_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_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)
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_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))