fix nets, improve drawing

This commit is contained in:
Sijmen 2019-03-08 15:40:50 +01:00
parent 8cdfa7b878
commit 9c467e1878

View file

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