400 lines
13 KiB
Python
400 lines
13 KiB
Python
import matplotlib.pyplot as plt
|
|
import matplotlib.colors
|
|
import numpy as np
|
|
import random
|
|
from collections import defaultdict
|
|
from enum import IntEnum
|
|
from sys import argv
|
|
from matplotlib.patches import Patch
|
|
|
|
|
|
class Model:
|
|
def __init__(self, width=50, height=50, humandens=0.15, mosquitodens=0.10,
|
|
immunepct=0.01, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
|
|
hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6,
|
|
mosqdiepct=10**-3, mosqnetdens=0.00, time_steps=50000,
|
|
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 mosquito to be infected by a human
|
|
self.hm_infpct = hm_infpct
|
|
# Chance for a human to be infected from a mosquito
|
|
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.infected_trend = np.zeros(time_steps)
|
|
self.immune_trend = np.zeros(time_steps)
|
|
|
|
self.grid = self.gen_humans()
|
|
self.mosquitos = self.gen_mosquitos()
|
|
self.nets = self.gen_nets()
|
|
|
|
# statistics
|
|
self.stats = defaultdict(int)
|
|
|
|
if self.graphical:
|
|
self.init_draw()
|
|
|
|
def init_draw(self):
|
|
plt.ion()
|
|
plt.rcParams.update({'font.size': 18})
|
|
self.colors = matplotlib.colors.ListedColormap(
|
|
["white", "green", "red", "yellow"])
|
|
|
|
def make_babies(self, n):
|
|
if n == 0:
|
|
return
|
|
|
|
self.stats["humans born"] += n
|
|
|
|
births = np.transpose(random.sample(
|
|
list(np.transpose(np.where(self.grid == Human.DEAD))), n))
|
|
|
|
self.grid[births[0], births[1]] = \
|
|
np.random.choice((Human.HEALTHY, Human.IMMUNE), size=n,
|
|
p=(1 - self.immunepct, self.immunepct))
|
|
|
|
# Randomly distribute a net
|
|
nets = births.T[np.random.rand(len(births.T)) < self.mosqnetdens].T
|
|
self.nets[nets[0], nets[1]] = True
|
|
|
|
def recycle_human(self):
|
|
"""
|
|
Determine if a human dies of natural causes and then replace them by a
|
|
new human.
|
|
"""
|
|
# Find living humans, determine if they die, and if so, kill them
|
|
humans = np.transpose(np.where(self.grid != Human.DEAD))
|
|
deaths = np.random.rand(len(humans)) < self.humandiepct
|
|
|
|
locations = humans[deaths].T
|
|
self.grid[locations[0], locations[1]] = Human.DEAD
|
|
self.nets[locations[0], locations[1]] = False
|
|
|
|
death_count = len(np.where(deaths)[0])
|
|
self.stats["natural deaths"] += death_count
|
|
|
|
# Replace the dead humans
|
|
self.make_babies(death_count)
|
|
|
|
def recycle_mosquito(self):
|
|
"""
|
|
Determine if a mosquito dies of natural causes and then replace it by
|
|
a new mosquito.
|
|
"""
|
|
indices = np.random.rand(len(self.mosquitos)) < self.mosqdiepct
|
|
deaths = self.mosquitos[indices]
|
|
for mosq in deaths:
|
|
mosq.hungry = False
|
|
mosq.infected = np.random.rand() < self.mosqinfpct
|
|
mosq.x = np.random.randint(0, self.width)
|
|
mosq.y = np.random.randint(0, self.height)
|
|
|
|
self.stats["mosquito deaths"] += len(deaths)
|
|
|
|
def do_malaria(self):
|
|
"""
|
|
This function determines who of the infected dies from their illness
|
|
"""
|
|
# Find infected humans, determine if they die, and if so, kill them
|
|
infected = np.transpose(np.where(self.grid == Human.INFECTED))
|
|
deaths = np.random.rand(len(infected)) < self.hinfdiepct
|
|
|
|
locs = infected[deaths].T
|
|
self.grid[locs[0], locs[1]] = Human.DEAD
|
|
self.nets[locs[0], locs[1]] = False
|
|
|
|
death_count = len(np.where(deaths)[0])
|
|
self.stats["malaria deaths"] += death_count
|
|
|
|
# Replace the dead humans
|
|
self.make_babies(death_count)
|
|
|
|
def feed(self):
|
|
"""
|
|
Feed the mosquitos that want to and can be fed
|
|
"""
|
|
# TODO: dit refactoren?
|
|
for mos in self.mosquitos:
|
|
if not mos.hungry:
|
|
continue
|
|
|
|
# state of current place on the grid where mosquito lives
|
|
state = self.grid[mos.x, mos.y]
|
|
|
|
if state == Human.DEAD:
|
|
continue
|
|
|
|
self.stats["mosquitos fed"] += 1
|
|
mos.hungry = False
|
|
|
|
# check if healthy human needs to be infected or mosquito
|
|
# becomes infected from eating
|
|
if state == Human.HEALTHY and mos.infected \
|
|
and random.uniform(0, 1) < self.mh_infpct:
|
|
self.grid[mos.x, mos.y] = Human.INFECTED
|
|
self.stats["humans infected"] += 1
|
|
elif state == Human.INFECTED and not mos.infected \
|
|
and random.uniform(0, 1) < self.hm_infpct:
|
|
mos.infected = True
|
|
self.stats["mosquitos infected"] += 1
|
|
|
|
def determine_hunger(self):
|
|
"""
|
|
Determines which mosquitos should get hungry
|
|
"""
|
|
for mos in self.mosquitos:
|
|
mos.hungry = not mos.hungry and \
|
|
random.uniform(0, 1) < self.mhungrypct
|
|
|
|
def get_movementbox(self, x: int, y: int):
|
|
"""
|
|
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)])[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 np.array(mosquitos)
|
|
|
|
def gen_nets(self):
|
|
"""
|
|
Generates the grid of nets
|
|
"""
|
|
humans = np.transpose(np.where(self.grid != Human.DEAD))
|
|
positions = humans[np.random.choice(
|
|
len(humans), size=round(self.mosqnetdens * len(humans)))].T
|
|
|
|
grid = np.zeros((self.width, self.height), dtype=bool)
|
|
grid[positions[0], positions[1]] = True
|
|
return grid
|
|
|
|
def run(self):
|
|
"""
|
|
This functions runs the simulation
|
|
"""
|
|
# Actual simulation runs inside try except to catch keyboard interrupts
|
|
# and always print stats
|
|
try:
|
|
for t in range(self.time_steps):
|
|
self.step()
|
|
|
|
total = np.count_nonzero(self.grid != Human.DEAD)
|
|
infected = np.count_nonzero(self.grid == Human.INFECTED) / \
|
|
total
|
|
immune = np.count_nonzero(self.grid == Human.IMMUNE) / total
|
|
|
|
self.infected_trend[t] = infected
|
|
self.immune_trend[t] = immune
|
|
print(f"t={t}, infected={infected:.2f}, immune={immune:.2f}",
|
|
end='\r')
|
|
|
|
if self.graphical and t % 1000 == 0:
|
|
self.draw(t)
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
print()
|
|
self.compile_stats()
|
|
self.print_stats()
|
|
self.draw_graph(t)
|
|
|
|
def compile_stats(self):
|
|
"""
|
|
Compiles a comprehensive list of statistics of the simulation
|
|
"""
|
|
self.stats["total deaths"] = \
|
|
self.stats["malaria deaths"] + self.stats["natural deaths"]
|
|
self.stats["final immunity"] = self.immune_trend[-1]
|
|
|
|
self.stats["net count"] = len(np.where(self.nets)[0])
|
|
|
|
def print_stats(self):
|
|
"""
|
|
Prints the gathered statistics from the simulation
|
|
"""
|
|
for stat, value in sorted(self.stats.items()):
|
|
print(f"{stat}: {self.stats[stat]}")
|
|
|
|
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()
|
|
# check if mosquitoes die
|
|
self.recycle_mosquito()
|
|
# move mosquitos
|
|
self.move_mosquitos()
|
|
# feed hungry mosquitos
|
|
self.feed()
|
|
# make mosquitos hungry again
|
|
self.determine_hunger()
|
|
|
|
def draw(self, t: int):
|
|
"""
|
|
Draws the grid of humans, tents and mosquitos
|
|
"""
|
|
plt.clf()
|
|
|
|
plt.title("t={}".format(t))
|
|
|
|
# draw the grid
|
|
plt.imshow(self.grid, cmap=self.colors)
|
|
|
|
# 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.y, mos.x, mos.get_color()+mos.get_shape())
|
|
|
|
# draw the legend
|
|
dead_patch = Patch(color="green", label="Healthy human")
|
|
immune_patch = Patch(color="yellow", label="Immune human")
|
|
infected_patch = Patch(color="red", label="Infected human")
|
|
plt.legend(handles=[dead_patch, immune_patch, infected_patch],
|
|
loc=9, bbox_to_anchor=(0.5, -0.03), ncol=5)
|
|
|
|
plt.pause(0.0001)
|
|
|
|
def draw_graph(self, t):
|
|
plt.ioff()
|
|
plt.clf()
|
|
|
|
plt.plot(np.arange(t), self.infected_trend[:t],
|
|
label="infected humans")
|
|
plt.plot(np.arange(t), self.immune_trend[:t],
|
|
label="immune humans")
|
|
|
|
plt.legend()
|
|
|
|
plt.xlabel("timestep")
|
|
plt.ylabel("prevalence")
|
|
|
|
plt.show()
|
|
|
|
|
|
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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
graphical = argv[1] != "-ng"
|
|
except IndexError:
|
|
graphical = True
|
|
|
|
model = Model(graphical=graphical)
|
|
model.run()
|
|
|