Add year 2

year2.py

@@ -0,0 +1,310 @@
+#!/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()