diff --git a/malaria.py b/malaria.py index c533453..71f6f08 100644 --- a/malaria.py +++ b/malaria.py @@ -11,7 +11,7 @@ class Model: hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6, mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000, graphical=True): - + self.width = width self.height = height @@ -59,17 +59,16 @@ class Model: if self.graphical: self.init_draw() - + def init_draw(self): plt.ion() 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): """ - Determine if a human dies of natural causes and then replace them by a new human + Determine if a human dies of natural causes and then replace them by a + new human. """ # Get all living humans humans = np.transpose(np.where(self.grid != Human.DEAD)) @@ -88,13 +87,15 @@ class Model: self.stats["natural deaths"] += death_count # Pick a random, unpopulated spot - births = np.array(random.sample(list(np.transpose(np.where(self.grid == Human.DEAD))), - death_count)) + 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]) + self.grid[birth[0]][birth[1]] = \ + np.random.choice((Human.HEALTHY, Human.IMMUNE), + p=(1 - self.immunepct, self.immunepct)) def do_malaria(self): """ @@ -112,10 +113,10 @@ class Model: self.stats["malaria deaths"] += len(np.where(deaths)[0]) def feed(self): - #TODO: dit refactoren? """ Feed the mosquitos that want to and can be fed """ + # TODO: dit refactoren? for mos in self.mosquitos: if not mos.hungry: continue @@ -125,12 +126,15 @@ class Model: if state != Human.DEAD: self.stats["mosquitos fed"] += 1 mos.hungry = False - - # check if healthy human needs to be infected or mosquito becomes infected from eating - if state == Human.HEALTHY and mos.infected and random.uniform(0, 1) < self.mh_infpct: + + # check if healthy human needs to be infected or mosquito + # becomes infected from eating + if state == Human.HEALTHY and mos.infected \ + and random.uniform(0, 1) < self.mh_infpct: self.grid[mos.x, mos.y] = Human.INFECTED self.stats["humans infected"] += 1 - elif state == Human.INFECTED and not mos.infected and random.uniform(0, 1) < self.hm_infpct: + elif state == Human.INFECTED and not mos.infected \ + and random.uniform(0, 1) < self.hm_infpct: self.stats["mosquitos infected"] += 1 mos.infected = True @@ -139,8 +143,9 @@ class Model: Determines which mosquitos should get hungry """ for mos in self.mosquitos: - mos.hungry = not mos.hungry and random.uniform(0, 1) < self.mhungrypct - + mos.hungry = not mos.hungry and \ + random.uniform(0, 1) < self.mhungrypct + def get_movementbox(self, x: int, y: int): """ Returns indices of a moore neighbourhood around the given index @@ -151,8 +156,9 @@ class Model: 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)] + 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)) @@ -168,8 +174,8 @@ class Model: movement = self.get_movementbox(mosq.x, mosq.y) # check for nets, and thus legal locations to go for the mosquito - legal_moves = np.where(self.nets[tuple(movement.T)] == False)[0] - + legal_moves = np.where(~self.nets[tuple(movement.T)])[0] + # choose random new position new_pos = random.choice(legal_moves) @@ -219,7 +225,7 @@ class Model: Generates the grid of nets """ - return np.random.choice([False, True], + return np.random.choice([False, True], p=[1-self.mosqnetdens, self.mosqnetdens], size=(self.width, self.height)) @@ -228,7 +234,8 @@ class Model: This functions runs the simulation """ print(chr(27) + "[2J") - # actual simulation runs inside try except to catch keyboard interrupts and always print stats + # Actual simulation runs inside try except to catch keyboard interrupts + # and always print stats try: for t in range(self.time_steps): print("Simulating timestep: {}".format(t), end='\r') @@ -237,7 +244,7 @@ class Model: self.draw(t) except KeyboardInterrupt: pass - + print(chr(27) + "[2J") self.compile_stats() self.print_stats() @@ -246,8 +253,9 @@ class Model: """ Compiles a comprehensive list of statistics of the simulation """ - self.stats["total deaths"] = self.stats["malaria deaths"] + self.stats["natural deaths"] - + self.stats["total deaths"] = \ + self.stats["malaria deaths"] + self.stats["natural deaths"] + # print(np.where(self.nets)) self.stats["net count"] = len(np.where(self.nets)[0]) @@ -272,7 +280,7 @@ class Model: self.feed() # make mosquitos hungry again self.determine_hunger() - + def draw(self, t: int): """ Draws the grid of humans, tents and mosquitos @@ -280,8 +288,8 @@ class Model: # this function draws the humans plt.title("t={}".format(t)) # draw the grid - plt.imshow(self.grid, cmap=self.colors, norm=self.norm) - + plt.imshow(self.grid, cmap=self.colors) + # draw nets net_locations = np.where(self.nets) plt.plot(net_locations[0], net_locations[1], 'w^') @@ -301,7 +309,7 @@ class Mosquito: 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"