class InfectionModel(Model): """ Mesa model class that simulates infection spread """ def __init__(self, params: dict): """ Parameters ---------- params: dict Simulation parameters """ self.current_id = 0 # inherited variable, for id generation self.params = params # parameters self.statistics = { # statistics for data collector "infected": 0, "recovered": 0, "susceptible": 0, "vaccinated": 0, "deaths": 0, "alive": 0, "total_infections": 0, "total_recoveries": 0, } self.grid = MultiGrid(self.params['grid_width'], self.params['grid_height'], True) # grid that agents move on self.schedule = SimultaneousActivation(self) # scheduler for iterations of the simulation self.dataCollector = DataCollector(model_reporters={ # to collect data for the graph "infected": lambda m: m.statistics["infected"], "recovered": lambda m: m.statistics["recovered"], "susceptible": lambda m: m.statistics["susceptible"], "vaccinated": lambda m: m.statistics["vaccinated"], "deaths": lambda m: m.statistics["deaths"], "alive": lambda m: m.statistics["alive"], "total_infections": lambda m: m.statistics["total_infections"], "total_recoveries": lambda m: m.statistics["total_recoveries"], }) self.running = True # required for visualization, tells if simulation is done self.dead_agents = [] # when agents die, they are added to this list to be removed self.step_count = 0 # number of steps completed, required for vaccination self.vaccination_started = False # has vaccination started? # creating agents for _ in range(self.params['num_agents']): # initial state of this agent initial_state = InfectionState.INF if \ self.random.uniform(0, 1) < self.params['initial_infected_chance'] else InfectionState.SUS # by default, this won't add to total_infections which leads to incorrect results if initial_state == InfectionState.INF: self.statistics["total_infections"] += 1 # randomise position pos = self.random.randrange(self.grid.width), self.random.randrange(self.grid.height) self.add_agent(self.create_agent(initial_state), pos) def check_running(self): """ Checks and returns if the simulation is still running (there are infected people) """ for agent in self.schedule.agent_buffer(): if agent.state == InfectionState.INF: return True return False def step(self): """ Called every step """ # just to show the progress while running in console if self.step_count % 100 == 0: print(self.step_count) self.per_agent_actions() # simulate actions to be taken globally on all agents self.schedule.step() # run step for all agents # collect data at a particular frequency if self.step_count % self.params['data_collection_frequency'] == 0: self.calculate_statistics() # calculate statistics for data collector self.dataCollector.collect(self) # collect data self.step_count += 1 # if vaccination is enabled and enough time has passed if not self.vaccination_started and self.params['vaccination_start'] != -1 and \ self.step_count > self.params['vaccination_start']: self.vaccination_started = True # start vaccination for x in self.dead_agents: # remove dead agents self.remove_agent(x) self.statistics["deaths"] += 1 # add to death count self.dead_agents = [] self.running = self.check_running() # is the simulation still running? def calculate_statistics(self): """ Calculates statistics each iteration, for more efficient data collection """ # reset all iteration specific parameters self.statistics["infected"] = 0 self.statistics["recovered"] = 0 self.statistics["susceptible"] = 0 self.statistics["vaccinated"] = 0 self.statistics["alive"] = 0 for agent in self.schedule.agent_buffer(): self.statistics["alive"] += 1 if agent.state == InfectionState.INF: self.statistics["infected"] += 1 elif agent.state == InfectionState.SUS: self.statistics["susceptible"] += 1 elif agent.state == InfectionState.REC: self.statistics["recovered"] += 1 elif agent.state == InfectionState.VAC: self.statistics["vaccinated"] += 1 def per_agent_actions(self): """ Simulates actions to be taken on a global scale per agent """ # this should only occur once a day if self.step_count % 24 != 0: return for agent in self.schedule.agent_buffer(): if self.random.uniform(0, 1) < self.params['external_infection_chance']: agent.state = InfectionState.INF self.statistics["total_infections"] += 1 def create_agent(self, initial_state: InfectionState) -> PersonAgent: """ Creates an agent, and returns it Parameters ---------- initial_state : InfectionState Initial infection state of this agent Returns ------- PersonAgent The agent created """ return PersonAgent(self.next_id(), self, initial_state) def add_agent(self, agent: Agent, pos: Tuple[int, int]): """ Adds an agent to the simulation Parameters ---------- agent : agent The agent to be added pos : Tuple[int, int] The position where this agent should be on the grid """ # add to scheduler self.schedule.add(agent) # assign position self.grid.place_agent(agent, pos) def remove_agent(self, agent: Agent): """ Removes an agent from the simulation Parameters ---------- agent : Agent The agent to be removed from the simulation """ self.grid.remove_agent(agent) self.schedule.remove(agent)