diff --git a/malaria.py b/malaria.py index ecdb1d1..f5af246 100644 --- a/malaria.py +++ b/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 @@ -87,8 +87,29 @@ 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() @@ -127,6 +184,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): @@ -138,5 +203,7 @@ class Human(IntEnum): if __name__ == "__main__": model = Model() - model.run() + # model.run() + + print(model.get_movementbox(0,0))