malariariariara

malaria.py

@@ -1,4 +1,5 @@
 import matplotlib.pyplot as plt
+import matplotlib.colors
 import numpy as np
 import random
 from enum import IntEnum
@@ -6,7 +7,7 @@ from enum import IntEnum
 
 class Model:
     def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
-                 immumepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
+                 immunepct=0.1, 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.05, time_steps=500):
         self.width = width
@@ -31,47 +32,50 @@ class Model:
 
     def recycle_human(self):
         # Get all living humans
-        humans = np.transpose(np.where(a != Human.DEAD))
-        deaths = numpy.random.rand(len(humans)) < self.humandiepct
-        grid[humans[deaths]] = Human.DEAD
+        humans = np.transpose(np.where(self.grid != Human.DEAD))
 
-        for x, y in humans:
-            if numpy.random.rand() < self.humandiepct:
-                grid[x, y] = 
+        # 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
 
         # Pick a random, unpopulated spot
-        x, y = random.choice(np.transpose(np.where(a == Human.DEAD)))
+        x, y = random.choice(np.transpose(np.where(self.grid == Human.DEAD)),
+                             size=np.count_nonzero(deaths))
 
     def gen_humans(self):
+        # Calculate the probabilities
         p_dead = 1 - self.humandens
         p_immune = self.humandens * self.immunepct
-        p_healthy = p_dead - p_immune
+        p_healthy = self.humandens - p_immune
 
+        # Create the grid with humans.
         return np.random.choice((Human.DEAD, Human.HEALTHY, Human.IMMUNE),
                                 size=(self.width, self.height),
-                                p=(p_dead, p_immune, p_healthy))
+                                p=(p_dead, p_healthy, p_immune))
 
     def gen_mosquitos(self):
-        mosquitos = []
-
         count = self.width * self.height * self.mosquitodens
-        for _ in count:
-
-
 
     def run(self):
         for _ in range(self.time_steps):
             self.step()
             self.draw()
 
-    
     def step(self):
         # dingen
         pass
 
     def draw(self):
         # this function draws the humans
-        plt.imshow(self.humans)
+        colors = matplotlib.colors.ListedColormap(
+            ["black", "green", "red", "yellow"])
+        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
+        norm = matplotlib.colors.BoundaryNorm(bounds, colors.N)
+
+        plt.imshow(self.humans, cmap=colors, norm=norm)
+
         plt.pause(0.0001)
         plt.clf()
 
@@ -93,11 +97,8 @@ class Human(IntEnum):
 
 
 if __name__ == "__main__":
-    
     plt.ion()
+    model = Model()
+    while True:
+        model.draw()
 
-    for i in range(500):
-        img = np.random.randint(0, 3, (10000, 10000))
-        plt.imshow(img)
-        plt.pause(0.0001)
-        plt.clf()