movementboxes added

malaria.py

@@ -21,7 +21,7 @@ class Model:
         self.mosquitodens = mosquitodens
         # Percentage of humans that are immune
         self.immunepct = immunepct
-        # Chance for a mosquito to be infected
+        # Chance for a mosquito to be infectuous
         self.mosqinfpct = mosqinfpct
         # Chance for a human to be infected by a mosquito bite
         self.hm_infpct = hm_infpct
@@ -88,7 +88,28 @@ class Model:
         # Now let's kill them
         self.grid[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD
     
+    def get_movementbox(self, x, y):
+        """
+        Returns indices of a moore neighbourhood around the given index
+        """
+
+        x_min = (x - 1)
+        x_max = (x + 1)
+        y_min = (y - 1)
+        y_max = (y + 1)
+
+        indices = [(i % self.width, j % self.height) for i in range(x_min, x_max + 1) 
+                                                     for j in range(y_min, y_max + 1)]
+
+        return indices
+
+    def move_mosquitos(self):
+        pass
+
     def gen_humans(self):
+        """
+        Fill the grid with humans that can either be healthy or infected
+        """
         # Calculate the probabilities
         p_dead = 1 - self.humandens
         p_immune = self.humandens * self.immunepct
@@ -100,21 +121,57 @@ class Model:
                                 p=(p_dead, p_healthy, p_immune))
 
     def gen_mosquitos(self):
-        count = self.width * self.height * self.mosquitodens
+        """
+        Generate the list of mosquitos
+        """
+
+        mosquitos = []
+
+        count = int(self.width * self.height * self.mosquitodens)
+
+        # generate random x and y coordinates
+        xs = np.random.randint(0, self.width, count)
+        ys = np.random.randint(0, self.height, count)
+        coords = list(zip(xs, ys))
+
+        # generate the mosquitos
+        for coord in coords:
+            # determine if the mosquito is infected
+            infected = random.uniform(0, 1) < self.mosqinfpct
+            # determine if the mosquito starts out hungry
+            hungry = random.uniform(0, 1) < self.mhungrypct
+            mosquitos.append(Mosquito(coord[0], coord[1], infected, hungry))
+
+        return mosquitos
 
     def run(self):
+        """
+        This functions runs the simulation
+        """
         for t in range(self.time_steps):
             self.step()
             self.draw(t)
 
     def step(self):
-        # self.do_malaria()
+        """
+        Step through a timestep of the simulation
+        """
+        # check who dies from malaria
+        self.do_malaria()
+        # check if people die from other causes
         self.recycle_human()
+        # move mosquitos
+        self.move_mosquitos()
         
-    def draw(self, t):
+
+    def draw(self, t: int):
         # this function draws the humans
         plt.title("t={}".format(t))
+        # draw the grid
         plt.imshow(self.grid, cmap=self.colors, norm=self.norm)
+        # draw mosquitos
+        for mos in self.mosquitos:
+            plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape())
 
         plt.pause(0.0001)
         plt.clf()
@@ -128,6 +185,14 @@ class Mosquito:
         self.infected = infected
         self.hungry = hungry
     
+    def get_color(self):
+        # returns the color for drawing, red if infected blue otherwise
+        return "r" if self.infected else "b"
+
+    def get_shape(self):
+        # return the shape for drawing, o if hungry + otherwise
+        return "o" if self.hungry else "+"
+
 
 class Human(IntEnum):
     DEAD = 0
@@ -138,5 +203,7 @@ class Human(IntEnum):
 
 if __name__ == "__main__":
     model = Model()
-    model.run()
+    # model.run()
+
+    print(model.get_movementbox(0,0))