diff --git a/malaria.py b/malaria.py index 8d6b4f9..9a2c334 100644 --- a/malaria.py +++ b/malaria.py @@ -44,7 +44,8 @@ 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()