From 9c467e187849b99e7c13c9b023b17872a3d88c00 Mon Sep 17 00:00:00 2001 From: Sijmen Schoon Date: Fri, 8 Mar 2019 15:40:50 +0100 Subject: [PATCH] fix nets, improve drawing --- malaria.py | 145 +++++++++++++++++++++++++++++++---------------------- 1 file changed, 85 insertions(+), 60 deletions(-) diff --git a/malaria.py b/malaria.py index 71f6f08..ef5961f 100644 --- a/malaria.py +++ b/malaria.py @@ -2,13 +2,16 @@ import matplotlib.pyplot as plt import matplotlib.colors import numpy as np import random +from collections import defaultdict from enum import IntEnum +from sys import argv +from matplotlib.patches import Patch 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**-6, + hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-5, mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000, graphical=True): @@ -47,15 +50,7 @@ class Model: self.nets = self.gen_nets() # statistics - self.stats = { - "natural deaths": 0, - "malaria deaths": 0, - "total deaths": 0, - "mosquitos fed": 0, - "humans infected": 0, - "mosquitos infected": 0, - "net count": 0 - } + self.stats = defaultdict(int) if self.graphical: self.init_draw() @@ -63,54 +58,61 @@ class Model: def init_draw(self): plt.ion() self.colors = matplotlib.colors.ListedColormap( - ["black", "green", "red", "yellow"]) + ["white", "green", "red", "yellow"]) + + def make_babies(self, n): + if n == 0: + return + + self.stats["humans born"] += n + + births = np.transpose(random.sample( + list(np.transpose(np.where(self.grid == Human.DEAD))), n)) + + self.grid[births[0], births[1]] = \ + np.random.choice((Human.HEALTHY, Human.IMMUNE), size=n, + p=(1 - self.immunepct, self.immunepct)) + + # Randomly distribute a net + nets = births.T[np.random.rand(len(births.T)) < self.mosqnetdens].T + self.nets[nets[0], nets[1]] = True def recycle_human(self): """ Determine if a human dies of natural causes and then replace them by a new human. """ - # Get all living humans + # Find living humans, determine if they die, and if so, kill them 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 + locations = humans[deaths].T + self.grid[locations[0], locations[1]] = Human.DEAD + self.nets[locations[0], locations[1]] = False + death_count = len(np.where(deaths)[0]) 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)) - - # 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)) + # Replace the dead humans + self.make_babies(death_count) def do_malaria(self): """ This function determines who of the infected dies from their illness """ - # Get all infected humans + # Find infected humans, determine if they die, and if so, kill them 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 + locs = infected[deaths].T + self.grid[locs[0], locs[1]] = Human.DEAD + self.nets[locs[0], locs[1]] = False - self.stats["malaria deaths"] += len(np.where(deaths)[0]) + death_count = len(np.where(deaths)[0]) + self.stats["malaria deaths"] += death_count + + # Replace the dead humans + self.make_babies(death_count) def feed(self): """ @@ -120,23 +122,26 @@ class Model: for mos in self.mosquitos: if not mos.hungry: continue + # state of current place on the grid where mosquito lives state = self.grid[mos.x, mos.y] - if state != Human.DEAD: - self.stats["mosquitos fed"] += 1 - mos.hungry = False + if state == Human.DEAD: + continue - # 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: - self.stats["mosquitos infected"] += 1 - mos.infected = True + 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: + 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: + mos.infected = True + self.stats["mosquitos infected"] += 1 def determine_hunger(self): """ @@ -224,18 +229,22 @@ class Model: """ Generates the grid of nets """ + humans = np.transpose(np.where(self.grid != Human.DEAD)) + positions = humans[np.random.choice( + len(humans), size=round(self.mosqnetdens * len(humans)))].T - return np.random.choice([False, True], - p=[1-self.mosqnetdens, self.mosqnetdens], - size=(self.width, self.height)) + grid = np.zeros((self.width, self.height), dtype=bool) + grid[positions[0], positions[1]] = True + return grid def run(self): """ This functions runs the simulation """ - print(chr(27) + "[2J") # Actual simulation runs inside try except to catch keyboard interrupts # and always print stats + self.stats["humans alive before simulation"] = \ + np.count_nonzero(self.grid != Human.DEAD) try: for t in range(self.time_steps): print("Simulating timestep: {}".format(t), end='\r') @@ -244,8 +253,10 @@ class Model: self.draw(t) except KeyboardInterrupt: pass + self.stats["humans alive after simulation"] = \ + np.count_nonzero(self.grid != Human.DEAD) - print(chr(27) + "[2J") + print() self.compile_stats() self.print_stats() @@ -256,14 +267,13 @@ class Model: 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]) def print_stats(self): """ Prints the gathered statistics from the simulation """ - for stat in self.stats: + for stat, value in sorted(self.stats.items()): print(f"{stat}: {self.stats[stat]}") def step(self): @@ -285,8 +295,11 @@ class Model: """ Draws the grid of humans, tents and mosquitos """ - # this function draws the humans + if t % 10 > 0: + return + plt.title("t={}".format(t)) + # draw the grid plt.imshow(self.grid, cmap=self.colors) @@ -296,7 +309,14 @@ class Model: # draw mosquitos for mos in self.mosquitos: - plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape()) + plt.plot(mos.y, mos.x, mos.get_color()+mos.get_shape()) + + # draw the legend + dead_patch = Patch(color="green", label="Healthy human") + immune_patch = Patch(color="yellow", label="Immune human") + infected_patch = Patch(color="red", label="Infected human") + plt.legend(handles=[dead_patch, immune_patch, infected_patch], + loc=9, bbox_to_anchor=(0.5, -0.03), ncol=5) plt.pause(0.0001) plt.clf() @@ -327,6 +347,11 @@ class Human(IntEnum): if __name__ == "__main__": - model = Model(graphical=True) + try: + graphical = argv[1] == "-g" + except IndexError: + graphical = False + + model = Model(graphical=graphical) model.run()