131 lines
4.1 KiB
Python
131 lines
4.1 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**-6,
|
|
mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000):
|
|
|
|
plt.ion()
|
|
self.width = width
|
|
self.height = height
|
|
|
|
# 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 infected
|
|
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.humans = self.gen_humans()
|
|
self.mosquitos = self.gen_mosquitos()
|
|
self.init_draw()
|
|
|
|
def init_draw(self):
|
|
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.humans != Human.DEAD))
|
|
|
|
# Get a mask of humans to kill
|
|
deaths = np.random.rand(len(humans)) < self.humandiepct
|
|
|
|
# Kill them.
|
|
self.humans[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD
|
|
|
|
# Pick a random, unpopulated spot
|
|
x, y = random.choice(np.transpose(np.where(self.humans == Human.DEAD)),
|
|
size=np.count_nonzero(deaths))
|
|
|
|
|
|
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.humans == Human.INFECTED))
|
|
|
|
# Decide which infected people die
|
|
deaths = np.random.rand(len(infected)) < self.hinfdiepct
|
|
|
|
# Now let's kill them
|
|
self.humans[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD
|
|
|
|
def gen_humans(self):
|
|
# 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):
|
|
count = self.width * self.height * self.mosquitodens
|
|
|
|
def run(self):
|
|
for t in range(self.time_steps):
|
|
self.step()
|
|
self.draw(t)
|
|
|
|
def step(self):
|
|
self.do_malaria()
|
|
|
|
def draw(self, t):
|
|
# this function draws the humans
|
|
plt.title("t={}".format(t))
|
|
plt.imshow(self.humans, cmap=self.colors, norm=self.norm)
|
|
|
|
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
|
|
|
|
|
|
class Human(IntEnum):
|
|
DEAD = 0
|
|
HEALTHY = 1
|
|
INFECTED = 2
|
|
IMMUNE = 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
model = Model()
|
|
model.run()
|
|
|