import matplotlib.pyplot as plt import matplotlib.colors import numpy as np import random from enum import IntEnum class Model: def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10, immunepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5, hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-3, mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000): plt.ion() self.width = width self.height = height # The percentage of tiles that start as humans self.humandens = humandens # The percentage of tiles that contain mosquitos self.mosquitodens = mosquitodens # Percentage of humans that are immune self.immunepct = immunepct # Chance for a mosquito to be infectuous self.mosqinfpct = mosqinfpct # Chance for a human to be infected by a mosquito bite self.hm_infpct = hm_infpct # Chance for a mosquito to be infected from biting an infected human self.mh_infpct = mh_infpct # Chance that an infected human dies self.hinfdiepct = hinfdiepct # Chance for a mosquito to be hungry self.mhungrypct = mhungrypct # Chance for human to die from random causes self.humandiepct = humandiepct # Chance for a mosquito to die self.mosqdiepct = mosqdiepct # Percentage of tiles that contain mosquito nets self.mosqnetdens = mosqnetdens # The number of timesteps to run te simulation for self.time_steps = time_steps self.grid = self.gen_humans() self.mosquitos = self.gen_mosquitos() self.init_draw() def init_draw(self): self.colors = matplotlib.colors.ListedColormap( ["black", "green", "red", "yellow"]) bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE] self.norm = matplotlib.colors.BoundaryNorm(bounds, self.colors.N) def recycle_human(self): # Get all living humans humans = np.transpose(np.where(self.grid != Human.DEAD)) # Get a mask of humans to kill deaths = np.random.rand(len(humans)) < self.humandiepct # Kill them. self.grid[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD # get num humans after killing humans_survive = len(np.transpose(np.where(self.grid != Human.DEAD))) death_count = len(humans) - humans_survive # Pick a random, unpopulated spot births = np.array(random.sample(list(np.transpose(np.where(self.grid == Human.DEAD))), death_count)) # Deliver the newborns for birth in births: self.grid[birth[0]][birth[1]] = np.random.choice([Human.HEALTHY, Human.IMMUNE], p=[1-self.immunepct, self.immunepct]) def do_malaria(self): """ This function determines who of the infected dies from their illness """ # Get all infected humans infected = np.transpose(np.where(self.grid == Human.INFECTED)) # Decide which infected people die deaths = np.random.rand(len(infected)) < self.hinfdiepct # Now let's kill them self.grid[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD def get_movementbox(self, x, y): """ Returns indices of a moore neighbourhood around the given index """ x_min = (x - 1) x_max = (x + 1) y_min = (y - 1) y_max = (y + 1) indices = [(i % self.width, j % self.height) for i in range(x_min, x_max + 1) for j in range(y_min, y_max + 1)] # remove current location from the indices indices.remove((x, y)) return indices def move_mosquitos(self): for mosq in self.mosquitos: # get the movement box for every mosquito movement = self.get_movementbox(mosq.x, mosq.y) # choose random new position new_pos = random.choice(movement) mosq.x = new_pos[0] mosq.y = new_pos[1] def gen_humans(self): """ Fill the grid with humans that can either be healthy or infected """ # Calculate the probabilities p_dead = 1 - self.humandens p_immune = self.humandens * self.immunepct p_healthy = self.humandens - p_immune # Create the grid with humans. return np.random.choice((Human.DEAD, Human.HEALTHY, Human.IMMUNE), size=(self.width, self.height), p=(p_dead, p_healthy, p_immune)) def gen_mosquitos(self): """ Generate the list of mosquitos """ mosquitos = [] count = int(self.width * self.height * self.mosquitodens) # generate random x and y coordinates xs = np.random.randint(0, self.width, count) ys = np.random.randint(0, self.height, count) coords = list(zip(xs, ys)) # generate the mosquitos for coord in coords: # determine if the mosquito is infected infected = random.uniform(0, 1) < self.mosqinfpct # determine if the mosquito starts out hungry hungry = random.uniform(0, 1) < self.mhungrypct mosquitos.append(Mosquito(coord[0], coord[1], infected, hungry)) return mosquitos def run(self): """ This functions runs the simulation """ for t in range(self.time_steps): self.step() self.draw(t) def step(self): """ Step through a timestep of the simulation """ # check who dies from malaria self.do_malaria() # check if people die from other causes self.recycle_human() # move mosquitos self.move_mosquitos() def draw(self, t: int): # this function draws the humans plt.title("t={}".format(t)) # draw the grid plt.imshow(self.grid, cmap=self.colors, norm=self.norm) # draw mosquitos for mos in self.mosquitos: plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape()) plt.pause(0.0001) plt.clf() class Mosquito: def __init__(self, x: int, y: int, infected: bool, hungry: bool): self.x = x self.y = y self.infected = infected self.hungry = hungry def get_color(self): # returns the color for drawing, red if infected blue otherwise return "r" if self.infected else "b" def get_shape(self): # return the shape for drawing, o if hungry + otherwise return "o" if self.hungry else "+" class Human(IntEnum): DEAD = 0 HEALTHY = 1 INFECTED = 2 IMMUNE = 3 if __name__ == "__main__": model = Model() model.run()