2019-ics-malariafreek/malaria.py

252 lines
7.9 KiB
Python

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