mosquitos now check for nets

malaria.py

@@ -44,6 +44,7 @@ class Model:
 
         self.grid = self.gen_humans()
         self.mosquitos = self.gen_mosquitos()
+        self.nets = self.gen_nets()
 
         if self.graphical:
             self.init_draw()
@@ -109,19 +110,24 @@ class Model:
         # remove current location from the indices
         indices.remove((x, y))
 
-        return indices
+        return np.array(indices)
 
     def move_mosquitos(self):
+        """
+        Move the mosquitos to a new location, checks for mosquito nets
+        """
         for mosq in self.mosquitos:
             # get the movement box for every mosquito
             movement = self.get_movementbox(mosq.x, mosq.y)
 
+            # check for nets, and thus legal locations to go for the mosquito
+            legal_moves = np.where(self.nets[tuple(movement.T)] == False)[0]
+            
             # choose random new position
-            new_pos = random.choice(movement)
-
-            mosq.x = new_pos[0]
-            mosq.y = new_pos[1]
+            new_pos = random.choice(legal_moves)
 
+            mosq.x = movement[new_pos][0]
+            mosq.y = movement[new_pos][1]
 
     def gen_humans(self):
         """
@@ -161,6 +167,16 @@ class Model:
 
         return mosquitos
 
+    def gen_nets(self):
+        """
+        Generates the grid of nets
+        """
+
+        return np.random.choice([False, True], 
+                                p=[1-self.mosqnetdens, self.mosqnetdens],
+                                size=(self.width, self.height))
+
+
     def run(self):
         """
         This functions runs the simulation
@@ -189,6 +205,11 @@ class Model:
         plt.title("t={}".format(t))
         # draw the grid
         plt.imshow(self.grid, cmap=self.colors, norm=self.norm)
+        
+        # draw nets
+        net_locations = np.where(self.nets)
+        plt.plot(net_locations[0], net_locations[1], 'w^')
+
         # draw mosquitos
         for mos in self.mosquitos:
             plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape())
@@ -220,8 +241,12 @@ class Human(IntEnum):
     INFECTED = 2
     IMMUNE = 3
 
+class Net(IntEnum):
+    NO_NET = 0
+    NET = 1
+
 
 if __name__ == "__main__":
-    model = Model(graphical=False)
+    model = Model(graphical=True)
     model.run()