2019-ics-malariafreek/malaria.py

import matplotlib.pyplot as plt
import matplotlib.colors
import numpy as np
import random
from enum import IntEnum


class Model:
    def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
                 immunepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
                 hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-3,
                 mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000,
                 graphical=True):
        
        self.width = width
        self.height = height

        # Determines if the simulation should be graphical
        self.graphical = graphical
        # The percentage of tiles that start as humans
        self.humandens = humandens
        # The percentage of tiles that contain mosquitos
        self.mosquitodens = mosquitodens
        # Percentage of humans that are immune
        self.immunepct = immunepct
        # 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
        # Chance for a mosquito to be infected from biting an infected human
        self.mh_infpct = mh_infpct
        # Chance that an infected human dies
        self.hinfdiepct = hinfdiepct
        # Chance for a mosquito to be hungry
        self.mhungrypct = mhungrypct
        # Chance for human to die from random causes
        self.humandiepct = humandiepct
        # Chance for a mosquito to die
        self.mosqdiepct = mosqdiepct
        # Percentage of tiles that contain mosquito nets
        self.mosqnetdens = mosqnetdens
        # The number of timesteps to run te simulation for
        self.time_steps = time_steps

        self.grid = self.gen_humans()
        self.mosquitos = self.gen_mosquitos()
        self.nets = self.gen_nets()

        if self.graphical:
            self.init_draw()
       
    def init_draw(self):
        plt.ion()
        self.colors = matplotlib.colors.ListedColormap(
            ["black", "green", "red", "yellow"])
        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
        self.norm = matplotlib.colors.BoundaryNorm(bounds, self.colors.N)

    def recycle_human(self):
        # Get all living humans
        humans = np.transpose(np.where(self.grid != Human.DEAD))

        # 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

        # get num humans after killing
        humans_survive = len(np.transpose(np.where(self.grid != Human.DEAD)))

        death_count = len(humans) - humans_survive


        # Pick a random, unpopulated spot
        births = np.array(random.sample(list(np.transpose(np.where(self.grid == Human.DEAD))),
                                        death_count))

        # Deliver the newborns
        for birth in births:
            self.grid[birth[0]][birth[1]] = np.random.choice([Human.HEALTHY, Human.IMMUNE], 
                                                                p=[1-self.immunepct, self.immunepct])

    def do_malaria(self):
        """
        This function determines who of the infected dies from their illness
        """
        # Get all infected humans
        infected = np.transpose(np.where(self.grid == Human.INFECTED))

        # Decide which infected people die
        deaths = np.random.rand(len(infected)) < self.hinfdiepct

        # 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)]

        # remove current location from the indices
        indices.remove((x, y))

        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(legal_moves)

            mosq.x = movement[new_pos][0]
            mosq.y = movement[new_pos][1]

    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
        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_healthy, p_immune))

    def gen_mosquitos(self):
        """
        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 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
        """
        for t in range(self.time_steps):
            self.step()
            if self.graphical:
                self.draw(t)
            else:
                print("Simulating timestep: {}".format(t), end='\r')

    def step(self):
        """
        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: 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 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())

        plt.pause(0.0001)
        plt.clf()


class Mosquito:
    def __init__(self, x: int, y: int, infected: bool, hungry: bool):
        self.x = x
        self.y = y

        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
    HEALTHY = 1
    INFECTED = 2
    IMMUNE = 3

class Net(IntEnum):
    NO_NET = 0
    NET = 1


if __name__ == "__main__":
    model = Model(graphical=True)
    model.run()