yeet

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()