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