diff --git a/malaria.py b/malaria.py index ef5961f..eed5f09 100644 --- a/malaria.py +++ b/malaria.py @@ -9,10 +9,10 @@ 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**-5, - mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000, + def __init__(self, width=50, height=50, humandens=0.15, mosquitodens=0.10, + immunepct=0.01, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5, + hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6, + mosqdiepct=10**-3, mosqnetdens=0.00, time_steps=50000, graphical=True): self.width = width @@ -45,6 +45,9 @@ class Model: # The number of timesteps to run te simulation for self.time_steps = time_steps + self.infected_trend = np.zeros(time_steps) + self.immune_trend = np.zeros(time_steps) + self.grid = self.gen_humans() self.mosquitos = self.gen_mosquitos() self.nets = self.gen_nets() @@ -57,6 +60,7 @@ class Model: def init_draw(self): plt.ion() + plt.rcParams.update({'font.size': 18}) self.colors = matplotlib.colors.ListedColormap( ["white", "green", "red", "yellow"]) @@ -96,6 +100,21 @@ class Model: # Replace the dead humans self.make_babies(death_count) + def recycle_mosquito(self): + """ + Determine if a mosquito dies of natural causes and then replace it by + a new mosquito. + """ + indices = np.random.rand(len(self.mosquitos)) < self.mosqdiepct + deaths = self.mosquitos[indices] + for mosq in deaths: + mosq.hungry = False + mosq.infected = np.random.rand() < self.mosqinfpct + mosq.x = np.random.randint(0, self.width) + mosq.y = np.random.randint(0, self.height) + + self.stats["mosquito deaths"] += len(deaths) + def do_malaria(self): """ This function determines who of the infected dies from their illness @@ -223,7 +242,7 @@ class Model: hungry = random.uniform(0, 1) < self.mhungrypct mosquitos.append(Mosquito(coord[0], coord[1], infected, hungry)) - return mosquitos + return np.array(mosquitos) def gen_nets(self): """ @@ -243,22 +262,29 @@ class Model: """ # 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') self.step() - if self.graphical: + + total = np.count_nonzero(self.grid != Human.DEAD) + infected = np.count_nonzero(self.grid == Human.INFECTED) / \ + total + immune = np.count_nonzero(self.grid == Human.IMMUNE) / total + + self.infected_trend[t] = infected + self.immune_trend[t] = immune + print(f"t={t}, infected={infected:.2f}, immune={immune:.2f}", + end='\r') + + if self.graphical and t % 1000 == 0: self.draw(t) except KeyboardInterrupt: pass - self.stats["humans alive after simulation"] = \ - np.count_nonzero(self.grid != Human.DEAD) print() self.compile_stats() self.print_stats() + self.draw_graph(t) def compile_stats(self): """ @@ -266,6 +292,7 @@ class Model: """ self.stats["total deaths"] = \ self.stats["malaria deaths"] + self.stats["natural deaths"] + self.stats["final immunity"] = self.immune_trend[-1] self.stats["net count"] = len(np.where(self.nets)[0]) @@ -284,6 +311,8 @@ class Model: self.do_malaria() # check if people die from other causes self.recycle_human() + # check if mosquitoes die + self.recycle_mosquito() # move mosquitos self.move_mosquitos() # feed hungry mosquitos @@ -295,8 +324,7 @@ class Model: """ Draws the grid of humans, tents and mosquitos """ - if t % 10 > 0: - return + plt.clf() plt.title("t={}".format(t)) @@ -319,8 +347,23 @@ class Model: loc=9, bbox_to_anchor=(0.5, -0.03), ncol=5) plt.pause(0.0001) + + def draw_graph(self, t): + plt.ioff() plt.clf() + plt.plot(np.arange(t), self.infected_trend[:t], + label="infected humans") + plt.plot(np.arange(t), self.immune_trend[:t], + label="immune humans") + + plt.legend() + + plt.xlabel("timestep") + plt.ylabel("prevalence") + + plt.show() + class Mosquito: def __init__(self, x: int, y: int, infected: bool, hungry: bool): @@ -348,9 +391,9 @@ class Human(IntEnum): if __name__ == "__main__": try: - graphical = argv[1] == "-g" + graphical = argv[1] != "-ng" except IndexError: - graphical = False + graphical = True model = Model(graphical=graphical) model.run()