fix nets, improve drawing

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 = {
+        self.stats = defaultdict(int)
-            "natural deaths": 0,
-            "malaria deaths": 0,
-            "total deaths": 0,
-            "mosquitos fed": 0,
-            "humans infected": 0,
-            "mosquitos infected": 0,
-            "net count": 0
-        }
 
         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.
+        locations = humans[deaths].T
-        self.grid[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD
+        self.grid[locations[0], locations[1]] = Human.DEAD
-
+        self.nets[locations[0], locations[1]] = False
-        # get num humans after killing
-        humans_survive = len(np.transpose(np.where(self.grid != Human.DEAD)))
-
-        death_count = len(humans) - humans_survive
 
+        death_count = len(np.where(deaths)[0])
         self.stats["natural deaths"] += death_count
 
-        # Pick a random, unpopulated spot
+        # Replace the dead humans
-        births = np.array(random.sample(
+        self.make_babies(death_count)
-            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
+        # 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
+        locs = infected[deaths].T
-        self.grid[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD
+        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,10 +122,13 @@ 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:
+            if state == Human.DEAD:
+                continue
+
             self.stats["mosquitos fed"] += 1
             mos.hungry = False
 
@@ -135,8 +140,8 @@ class Model:
                 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 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],
+        grid = np.zeros((self.width, self.height), dtype=bool)
-                                p=[1-self.mosqnetdens, self.mosqnetdens],
+        grid[positions[0], positions[1]] = True
-                                size=(self.width, self.height))
+        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()