fix nets, improve drawing
This commit is contained in:
parent
8cdfa7b878
commit
9c467e1878
1 changed files with 85 additions and 60 deletions
145
malaria.py
145
malaria.py
|
@ -2,13 +2,16 @@ import matplotlib.pyplot as plt
|
||||||
import matplotlib.colors
|
import matplotlib.colors
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import random
|
import random
|
||||||
|
from collections import defaultdict
|
||||||
from enum import IntEnum
|
from enum import IntEnum
|
||||||
|
from sys import argv
|
||||||
|
from matplotlib.patches import Patch
|
||||||
|
|
||||||
|
|
||||||
class Model:
|
class Model:
|
||||||
def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
|
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,
|
immunepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
|
||||||
hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6,
|
hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-5,
|
||||||
mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000,
|
mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000,
|
||||||
graphical=True):
|
graphical=True):
|
||||||
|
|
||||||
|
@ -47,15 +50,7 @@ class Model:
|
||||||
self.nets = self.gen_nets()
|
self.nets = self.gen_nets()
|
||||||
|
|
||||||
# statistics
|
# statistics
|
||||||
self.stats = {
|
self.stats = defaultdict(int)
|
||||||
"natural deaths": 0,
|
|
||||||
"malaria deaths": 0,
|
|
||||||
"total deaths": 0,
|
|
||||||
"mosquitos fed": 0,
|
|
||||||
"humans infected": 0,
|
|
||||||
"mosquitos infected": 0,
|
|
||||||
"net count": 0
|
|
||||||
}
|
|
||||||
|
|
||||||
if self.graphical:
|
if self.graphical:
|
||||||
self.init_draw()
|
self.init_draw()
|
||||||
|
@ -63,54 +58,61 @@ class Model:
|
||||||
def init_draw(self):
|
def init_draw(self):
|
||||||
plt.ion()
|
plt.ion()
|
||||||
self.colors = matplotlib.colors.ListedColormap(
|
self.colors = matplotlib.colors.ListedColormap(
|
||||||
["black", "green", "red", "yellow"])
|
["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):
|
def recycle_human(self):
|
||||||
"""
|
"""
|
||||||
Determine if a human dies of natural causes and then replace them by a
|
Determine if a human dies of natural causes and then replace them by a
|
||||||
new human.
|
new human.
|
||||||
"""
|
"""
|
||||||
# Get all living humans
|
# Find living humans, determine if they die, and if so, kill them
|
||||||
humans = np.transpose(np.where(self.grid != Human.DEAD))
|
humans = np.transpose(np.where(self.grid != Human.DEAD))
|
||||||
|
|
||||||
# Get a mask of humans to kill
|
|
||||||
deaths = np.random.rand(len(humans)) < self.humandiepct
|
deaths = np.random.rand(len(humans)) < self.humandiepct
|
||||||
|
|
||||||
# Kill them.
|
locations = humans[deaths].T
|
||||||
self.grid[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD
|
self.grid[locations[0], locations[1]] = Human.DEAD
|
||||||
|
self.nets[locations[0], locations[1]] = False
|
||||||
# get num humans after killing
|
|
||||||
humans_survive = len(np.transpose(np.where(self.grid != Human.DEAD)))
|
|
||||||
|
|
||||||
death_count = len(humans) - humans_survive
|
|
||||||
|
|
||||||
|
death_count = len(np.where(deaths)[0])
|
||||||
self.stats["natural deaths"] += death_count
|
self.stats["natural deaths"] += death_count
|
||||||
|
|
||||||
# Pick a random, unpopulated spot
|
# Replace the dead humans
|
||||||
births = np.array(random.sample(
|
self.make_babies(death_count)
|
||||||
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):
|
def do_malaria(self):
|
||||||
"""
|
"""
|
||||||
This function determines who of the infected dies from their illness
|
This function determines who of the infected dies from their illness
|
||||||
"""
|
"""
|
||||||
# Get all infected humans
|
# Find infected humans, determine if they die, and if so, kill them
|
||||||
infected = np.transpose(np.where(self.grid == Human.INFECTED))
|
infected = np.transpose(np.where(self.grid == Human.INFECTED))
|
||||||
|
|
||||||
# Decide which infected people die
|
|
||||||
deaths = np.random.rand(len(infected)) < self.hinfdiepct
|
deaths = np.random.rand(len(infected)) < self.hinfdiepct
|
||||||
|
|
||||||
# Now let's kill them
|
locs = infected[deaths].T
|
||||||
self.grid[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD
|
self.grid[locs[0], locs[1]] = Human.DEAD
|
||||||
|
self.nets[locs[0], locs[1]] = False
|
||||||
|
|
||||||
self.stats["malaria deaths"] += len(np.where(deaths)[0])
|
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):
|
def feed(self):
|
||||||
"""
|
"""
|
||||||
|
@ -120,23 +122,26 @@ class Model:
|
||||||
for mos in self.mosquitos:
|
for mos in self.mosquitos:
|
||||||
if not mos.hungry:
|
if not mos.hungry:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# state of current place on the grid where mosquito lives
|
# state of current place on the grid where mosquito lives
|
||||||
state = self.grid[mos.x, mos.y]
|
state = self.grid[mos.x, mos.y]
|
||||||
|
|
||||||
if state != Human.DEAD:
|
if state == Human.DEAD:
|
||||||
self.stats["mosquitos fed"] += 1
|
continue
|
||||||
mos.hungry = False
|
|
||||||
|
|
||||||
# check if healthy human needs to be infected or mosquito
|
self.stats["mosquitos fed"] += 1
|
||||||
# becomes infected from eating
|
mos.hungry = False
|
||||||
if state == Human.HEALTHY and mos.infected \
|
|
||||||
and random.uniform(0, 1) < self.mh_infpct:
|
# check if healthy human needs to be infected or mosquito
|
||||||
self.grid[mos.x, mos.y] = Human.INFECTED
|
# becomes infected from eating
|
||||||
self.stats["humans infected"] += 1
|
if state == Human.HEALTHY and mos.infected \
|
||||||
elif state == Human.INFECTED and not mos.infected \
|
and random.uniform(0, 1) < self.mh_infpct:
|
||||||
and random.uniform(0, 1) < self.hm_infpct:
|
self.grid[mos.x, mos.y] = Human.INFECTED
|
||||||
self.stats["mosquitos infected"] += 1
|
self.stats["humans infected"] += 1
|
||||||
mos.infected = True
|
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):
|
def determine_hunger(self):
|
||||||
"""
|
"""
|
||||||
|
@ -224,18 +229,22 @@ class Model:
|
||||||
"""
|
"""
|
||||||
Generates the grid of nets
|
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
|
||||||
|
|
||||||
return np.random.choice([False, True],
|
grid = np.zeros((self.width, self.height), dtype=bool)
|
||||||
p=[1-self.mosqnetdens, self.mosqnetdens],
|
grid[positions[0], positions[1]] = True
|
||||||
size=(self.width, self.height))
|
return grid
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
"""
|
"""
|
||||||
This functions runs the simulation
|
This functions runs the simulation
|
||||||
"""
|
"""
|
||||||
print(chr(27) + "[2J")
|
|
||||||
# Actual simulation runs inside try except to catch keyboard interrupts
|
# Actual simulation runs inside try except to catch keyboard interrupts
|
||||||
# and always print stats
|
# and always print stats
|
||||||
|
self.stats["humans alive before simulation"] = \
|
||||||
|
np.count_nonzero(self.grid != Human.DEAD)
|
||||||
try:
|
try:
|
||||||
for t in range(self.time_steps):
|
for t in range(self.time_steps):
|
||||||
print("Simulating timestep: {}".format(t), end='\r')
|
print("Simulating timestep: {}".format(t), end='\r')
|
||||||
|
@ -244,8 +253,10 @@ class Model:
|
||||||
self.draw(t)
|
self.draw(t)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
pass
|
pass
|
||||||
|
self.stats["humans alive after simulation"] = \
|
||||||
|
np.count_nonzero(self.grid != Human.DEAD)
|
||||||
|
|
||||||
print(chr(27) + "[2J")
|
print()
|
||||||
self.compile_stats()
|
self.compile_stats()
|
||||||
self.print_stats()
|
self.print_stats()
|
||||||
|
|
||||||
|
@ -256,14 +267,13 @@ class Model:
|
||||||
self.stats["total deaths"] = \
|
self.stats["total deaths"] = \
|
||||||
self.stats["malaria deaths"] + self.stats["natural deaths"]
|
self.stats["malaria deaths"] + self.stats["natural deaths"]
|
||||||
|
|
||||||
# print(np.where(self.nets))
|
|
||||||
self.stats["net count"] = len(np.where(self.nets)[0])
|
self.stats["net count"] = len(np.where(self.nets)[0])
|
||||||
|
|
||||||
def print_stats(self):
|
def print_stats(self):
|
||||||
"""
|
"""
|
||||||
Prints the gathered statistics from the simulation
|
Prints the gathered statistics from the simulation
|
||||||
"""
|
"""
|
||||||
for stat in self.stats:
|
for stat, value in sorted(self.stats.items()):
|
||||||
print(f"{stat}: {self.stats[stat]}")
|
print(f"{stat}: {self.stats[stat]}")
|
||||||
|
|
||||||
def step(self):
|
def step(self):
|
||||||
|
@ -285,8 +295,11 @@ class Model:
|
||||||
"""
|
"""
|
||||||
Draws the grid of humans, tents and mosquitos
|
Draws the grid of humans, tents and mosquitos
|
||||||
"""
|
"""
|
||||||
# this function draws the humans
|
if t % 10 > 0:
|
||||||
|
return
|
||||||
|
|
||||||
plt.title("t={}".format(t))
|
plt.title("t={}".format(t))
|
||||||
|
|
||||||
# draw the grid
|
# draw the grid
|
||||||
plt.imshow(self.grid, cmap=self.colors)
|
plt.imshow(self.grid, cmap=self.colors)
|
||||||
|
|
||||||
|
@ -296,7 +309,14 @@ class Model:
|
||||||
|
|
||||||
# draw mosquitos
|
# draw mosquitos
|
||||||
for mos in self.mosquitos:
|
for mos in self.mosquitos:
|
||||||
plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape())
|
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)
|
plt.pause(0.0001)
|
||||||
plt.clf()
|
plt.clf()
|
||||||
|
@ -327,6 +347,11 @@ class Human(IntEnum):
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
model = Model(graphical=True)
|
try:
|
||||||
|
graphical = argv[1] == "-g"
|
||||||
|
except IndexError:
|
||||||
|
graphical = False
|
||||||
|
|
||||||
|
model = Model(graphical=graphical)
|
||||||
model.run()
|
model.run()
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue