311 lines
10 KiB
Python
311 lines
10 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
import random
|
||
|
|
import time
|
||
|
|
from matplotlib import pyplot
|
||
|
|
|
||
|
|
HUMAN_DEAD = 0
|
||
|
|
HUMAN_HEALTHY = 1
|
||
|
|
HUMAN_INFECTED = 2
|
||
|
|
HUMAN_IMMUNE = 3
|
||
|
|
|
||
|
|
MOSQ_POSITION = 0
|
||
|
|
MOSQ_INFECTED = 1
|
||
|
|
MOSQ_HUNGRY = 2
|
||
|
|
|
||
|
|
|
||
|
|
class Model:
|
||
|
|
def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
|
||
|
|
immumepct=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):
|
||
|
|
self.width = width
|
||
|
|
self.height = height
|
||
|
|
|
||
|
|
# The initial density of humans
|
||
|
|
self.humandens = humandens
|
||
|
|
|
||
|
|
# The initial density mosquitos
|
||
|
|
self.mosquitodens = mosquitodens
|
||
|
|
|
||
|
|
# The chance that a human is immune
|
||
|
|
self.immumepct = immumepct
|
||
|
|
|
||
|
|
# The chance that a mosquito is born infected
|
||
|
|
self.mosqinfpct = mosqinfpct
|
||
|
|
|
||
|
|
# The chance that a human infects a mosquito
|
||
|
|
self.hm_infpct = hm_infpct
|
||
|
|
|
||
|
|
# The chance that a mosquito infects a human
|
||
|
|
self.mh_infpct = mh_infpct
|
||
|
|
|
||
|
|
# The chance that a mosquito gets hungry
|
||
|
|
self.mhungrypct = mhungrypct
|
||
|
|
|
||
|
|
# The chance that an infected human dies
|
||
|
|
self.hinfdiepct = hinfdiepct
|
||
|
|
|
||
|
|
# The chance that an uninfected/immume human dies
|
||
|
|
self.humandiepct = humandiepct
|
||
|
|
|
||
|
|
# The chance that a mosquito dies
|
||
|
|
self.mosqdiepct = mosqdiepct
|
||
|
|
|
||
|
|
# The density of mosquito nets
|
||
|
|
self.mosqnetdens = mosqnetdens
|
||
|
|
|
||
|
|
maxhumans = round(self.humandens * self.width * self.height)
|
||
|
|
maxmosquitos = round(self.mosquitodens * self.width * self.height)
|
||
|
|
maxmosqnets = round(self.mosqnetdens * self.width * self.height)
|
||
|
|
|
||
|
|
# Human: 0=dead, 1=healthy, 2=infected, 3=immume
|
||
|
|
# Mosquito: [(x, y), infected, hungry]
|
||
|
|
self.humans = [[HUMAN_DEAD for _ in range(width)]
|
||
|
|
for _ in range(height)]
|
||
|
|
self.mosquitos = []
|
||
|
|
self.mosqnets = [[False]*width for _ in range(height)]
|
||
|
|
for _ in range(maxhumans):
|
||
|
|
self.addhuman()
|
||
|
|
for _ in range(maxmosquitos):
|
||
|
|
self.addmosquito(random.uniform(0, 1) < mosqinfpct)
|
||
|
|
for _ in range(maxmosqnets):
|
||
|
|
self.addmosqnet()
|
||
|
|
|
||
|
|
def addhuman(self):
|
||
|
|
x, y = random.choice([(x, y)
|
||
|
|
for x in range(self.width) for y in range(self.height)
|
||
|
|
if self.humans[y][x] == HUMAN_DEAD])
|
||
|
|
|
||
|
|
if random.uniform(0, 1) < self.immumepct:
|
||
|
|
self.humans[y][x] = HUMAN_IMMUNE
|
||
|
|
else:
|
||
|
|
self.humans[y][x] = HUMAN_HEALTHY
|
||
|
|
|
||
|
|
def addmosquito(self, infected=False):
|
||
|
|
height = len(self.humans)
|
||
|
|
width = len(self.humans[0])
|
||
|
|
self.mosquitos.append([(random.randint(0, width-1),
|
||
|
|
random.randint(0, height-1)), infected, False])
|
||
|
|
|
||
|
|
def addmosqnet(self):
|
||
|
|
x, y = random.choice([(x, y)
|
||
|
|
for x in range(self.width) for y in range(self.height)
|
||
|
|
if self.mosqnets[y][x] == False])
|
||
|
|
self.mosqnets[y][x] = True
|
||
|
|
|
||
|
|
def move(self, pos):
|
||
|
|
x, y = pos
|
||
|
|
possibs = [(x, y) for x, y
|
||
|
|
in [(x-1, y), (x, y-1), (x+1, y), (x, y+1)]
|
||
|
|
if x > -1 and x < self.width
|
||
|
|
and y > -1 and y < self.height
|
||
|
|
and not self.mosqnets[y][x]]
|
||
|
|
try:
|
||
|
|
return random.choice(possibs)
|
||
|
|
except IndexError:
|
||
|
|
return (x, y)
|
||
|
|
|
||
|
|
def getstates(self):
|
||
|
|
humanhealthy = sum(1 for row in self.humans for human in row if human == HUMAN_HEALTHY)
|
||
|
|
humaninfected = sum(1 for row in self.humans for human in row if human == HUMAN_INFECTED)
|
||
|
|
humanimmune = sum(1 for row in self.humans for human in row if human == HUMAN_IMMUNE)
|
||
|
|
mosqinfected = sum(1 for mosq in self.mosquitos if mosq[MOSQ_INFECTED])
|
||
|
|
mosqhungry = sum(1 for mosq in self.mosquitos if mosq[MOSQ_HUNGRY])
|
||
|
|
return humanhealthy, humaninfected, humanimmune, mosqinfected, mosqhungry
|
||
|
|
|
||
|
|
def step(self):
|
||
|
|
# Copy the human values
|
||
|
|
humans = [human_row[:] for human_row in self.humans]
|
||
|
|
|
||
|
|
# Now I am become death, the destroyer of worlds
|
||
|
|
killed = 0
|
||
|
|
for human_row in humans:
|
||
|
|
for i, human in enumerate(human_row):
|
||
|
|
if human == HUMAN_INFECTED \
|
||
|
|
and random.uniform(0, 1) < self.hinfdiepct:
|
||
|
|
human_row[i] = HUMAN_DEAD
|
||
|
|
killed += 1
|
||
|
|
elif human != HUMAN_DEAD \
|
||
|
|
and random.uniform(0, 1) < self.humandiepct:
|
||
|
|
human_row[i] = HUMAN_DEAD
|
||
|
|
killed += 1
|
||
|
|
|
||
|
|
mosqkilled = 0
|
||
|
|
mosquitos = []
|
||
|
|
for pos, infected, hungry in self.mosquitos:
|
||
|
|
# Randomly move the mosquito to a new position
|
||
|
|
(x, y) = self.move(pos)
|
||
|
|
|
||
|
|
if hungry and self.humans[y][x]:
|
||
|
|
# If a mosquito is hungry and in the same cell as a mosquito,
|
||
|
|
# human, it'll bite
|
||
|
|
hungry = False
|
||
|
|
human = self.humans[y][x]
|
||
|
|
if human == HUMAN_INFECTED:
|
||
|
|
# If the human is infected, become infected
|
||
|
|
infected = random.uniform(0, 1) < self.hm_infpct
|
||
|
|
|
||
|
|
if infected and human == HUMAN_HEALTHY and \
|
||
|
|
random.uniform(0, 1) < self.mh_infpct:
|
||
|
|
# If the human is healthy and not immune, there's a
|
||
|
|
# mh_infpct chance it will become infected.
|
||
|
|
humans[y][x] = HUMAN_INFECTED
|
||
|
|
|
||
|
|
elif not hungry:
|
||
|
|
# If the mosquito is not hungry, there's a mhungrypct
|
||
|
|
# chance it'll become hungry
|
||
|
|
hungry = random.uniform(0, 1) < self.mhungrypct
|
||
|
|
|
||
|
|
# There's a mosqdiepct% chance that a mosq will die
|
||
|
|
if random.uniform(0, 1) < self.mosqdiepct:
|
||
|
|
mosqkilled += 1
|
||
|
|
else:
|
||
|
|
mosquitos.append(((x, y), infected, hungry))
|
||
|
|
|
||
|
|
# Send all new values to the model
|
||
|
|
self.mosquitos = mosquitos
|
||
|
|
self.humans = humans
|
||
|
|
|
||
|
|
for _ in range(killed):
|
||
|
|
# Add a human for each one that had died
|
||
|
|
self.addhuman()
|
||
|
|
for _ in range(mosqkilled):
|
||
|
|
# Add a mosqito for each one that had died
|
||
|
|
self.addmosquito()
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def drawhuman(human):
|
||
|
|
if human == HUMAN_DEAD:
|
||
|
|
return (31, 31, 31)
|
||
|
|
elif human == HUMAN_HEALTHY:
|
||
|
|
return (0, 127, 0)
|
||
|
|
elif human == HUMAN_INFECTED:
|
||
|
|
return (255, 0, 0)
|
||
|
|
elif human == HUMAN_IMMUNE:
|
||
|
|
return (0, 255, 0)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def drawmosq(mosq):
|
||
|
|
_, hungry, infected = mosq
|
||
|
|
if hungry and infected:
|
||
|
|
return (255, 61, 0)
|
||
|
|
elif hungry: # and not infected
|
||
|
|
return 191, 128, 0
|
||
|
|
elif infected: # and not hungry
|
||
|
|
return (127, 0, 0)
|
||
|
|
else: # neither
|
||
|
|
return (0, 0, 223)
|
||
|
|
|
||
|
|
def printlegend(self):
|
||
|
|
print('Humans (large blocks): ', end='')
|
||
|
|
print('\033[38;2;%i;%i;%imDead\033[0m' %
|
||
|
|
self.drawhuman(HUMAN_DEAD), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imHealthy\033[0m' %
|
||
|
|
self.drawhuman(HUMAN_HEALTHY), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imInfected\033[0m' %
|
||
|
|
self.drawhuman(HUMAN_INFECTED), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imImmune\033[0m' %
|
||
|
|
self.drawhuman(HUMAN_IMMUNE))
|
||
|
|
|
||
|
|
print('Mosquitos (small blocks): ', end='')
|
||
|
|
print('\033[38;2;%i;%i;%imHungry and Infected\033[0m' %
|
||
|
|
self.drawmosq(((0, 0), True, True)), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imHungry\033[0m' %
|
||
|
|
self.drawmosq(((0, 0), True, False)), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imInfected\033[0m' %
|
||
|
|
self.drawmosq(((0, 0), False, True)), end=', ')
|
||
|
|
print('\033[38;2;%i;%i;%imNeither\033[0m' %
|
||
|
|
self.drawmosq(((0, 0), False, False)))
|
||
|
|
|
||
|
|
print('Mosquito nets: /\\')
|
||
|
|
|
||
|
|
def print(self):
|
||
|
|
data = [[self.drawhuman(human) for human in row]
|
||
|
|
for row in self.humans]
|
||
|
|
|
||
|
|
# Draw the mosquitos
|
||
|
|
mosqs = [[None] * self.width for _ in range(self.height)]
|
||
|
|
for mosq in self.mosquitos:
|
||
|
|
x, y = mosq[0]
|
||
|
|
mosqs[y][x] = self.drawmosq(mosq)
|
||
|
|
|
||
|
|
print('\033[2J\033[0;0H' + '\033[0m\n'.join(''.join([
|
||
|
|
'\033[48;2;%i;%i;%im' % data[row][x] +
|
||
|
|
('\033[38;2;%i;%i;%im\u2590\u258c' % mosqs[row][x] if mosqs[row][x]
|
||
|
|
else '\033[38;2;255;255;255m/\\' if self.mosqnets[row][x]
|
||
|
|
else ' ')
|
||
|
|
for x in range(self.width)])
|
||
|
|
for row in range(0, self.height)) + '\033[0m')
|
||
|
|
|
||
|
|
self.printlegend()
|
||
|
|
|
||
|
|
def gui(self, printsteps=1):
|
||
|
|
try:
|
||
|
|
fps = 15
|
||
|
|
while True:
|
||
|
|
prevtime = time.time()
|
||
|
|
for _ in range(printsteps):
|
||
|
|
self.step()
|
||
|
|
self.print()
|
||
|
|
time.sleep(max(0, (1/fps) - (time.time()-prevtime)))
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
def frange(start, end, steps):
|
||
|
|
stepsize = (end - start) / (steps-1)
|
||
|
|
for n in range(steps):
|
||
|
|
yield start + stepsize * n
|
||
|
|
|
||
|
|
|
||
|
|
def statesvstime(tmax=40000):
|
||
|
|
model = Model(64, 64)
|
||
|
|
tr = list(range(tmax))
|
||
|
|
results = [None]*tmax
|
||
|
|
for t in tr:
|
||
|
|
print(t/tmax, end='\r')
|
||
|
|
model.step()
|
||
|
|
results[t] = model.getstates()
|
||
|
|
humanhealthy, humaninfected, humanimmune, mosqinfected, mosqhungry = zip(*results)
|
||
|
|
|
||
|
|
pyplot.plot(tr, humanhealthy, label='Healthy humans')
|
||
|
|
pyplot.plot(tr, humaninfected, label='Infected humans')
|
||
|
|
pyplot.plot(tr, humanimmune, label='Immune humans')
|
||
|
|
pyplot.plot(tr, mosqinfected, label='Infected mosquitos')
|
||
|
|
pyplot.plot(tr, mosqhungry, label='Hungry Mosquitos')
|
||
|
|
pyplot.legend()
|
||
|
|
pyplot.xlabel('time (steps)')
|
||
|
|
pyplot.ylabel('amount')
|
||
|
|
pyplot.show()
|
||
|
|
|
||
|
|
|
||
|
|
def netvsimmune(tmin=10000, denscnt=200):
|
||
|
|
dr = list(frange(0, 1, denscnt))
|
||
|
|
results = [None]*denscnt
|
||
|
|
for i, dens in enumerate(dr):
|
||
|
|
print(dens, end='\r')
|
||
|
|
model = Model(16, 16, mosqnetdens=dens)
|
||
|
|
for t in range(tmin):
|
||
|
|
model.step()
|
||
|
|
results[i] = sum(1 for row in model.humans for human in row if human == HUMAN_IMMUNE)
|
||
|
|
|
||
|
|
pyplot.plot(dr, results, label='Immunity')
|
||
|
|
pyplot.legend()
|
||
|
|
pyplot.xlabel('Mosquito net density')
|
||
|
|
pyplot.ylabel('Immune humans')
|
||
|
|
pyplot.show()
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
model = Model()
|
||
|
|
model.gui()
|
||
|
|
|
||
|
|
statesvstime()
|
||
|
|
netvsimmune()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
random.seed()
|
||
|
|
main()
|