Added code files

PandemicModel.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri May 15 19:08:07 2020
+
+@author: ameanasad
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from sklearn.neighbors import KDTree
+import random as r
+plt.rcParams['animation.ffmpeg_path'] = '/usr/local/bin/ffmpeg'
+
+# Random Seed set to 0 initializes the same pseudorandom number generator for multiple experiments.
+# Comment the random seeds out to use a different random generator everytime you run the code
+r.seed(0)
+np.random.seed(0)
+
+class PandemicSpread(object):
+
+    def __init__(self, size, nodes, infectionRate, radius, speed):
+        #nodes is the number of particle in the experiment
+        self.nodeSize = nodes
+        # Size is the dimension of the enclosing container for the particles
+        self.size = size
+        self.nodes_all = {"Infected": [],
+                      "Susceptible": [],
+                      "Dead": [],
+                      "Recovered" : []}
+        self.infectionRate = infectionRate
+        self.radius = radius
+        self.speed = speed
+
+    def deployNodes(self):
+        """
+        Deploy all the nodes in random positions
+        """
+        for i in range(self.nodeSize):
+            newNode = Node(self.size, "Susceptible", self.speed)
+            self.nodes_all['Susceptible'].append(newNode)   
+
+    def startInfection(self):
+        """
+        Start infection by deploying one random infected node. 
+        """
+        newNode = Node(self.size, "Infected", self.speed )
+        self.nodes_all['Infected'].append(newNode)
+
+    def checkInfection(self, node):
+        """
+        
+
+        Parameters
+        ----------
+        node : Object
+            A susceptible type node object.
+
+        Returns
+        -------
+        Boolean
+            True if node becomes infected, false otherwise.
+
+        """
+
+        # Used a K-dimensional tree to find neighbors within a givin radius
+        infected = self.getInfectedCoordinates()
+        infected = np.array(infected)
+        node = np.array([node.getCoordinates()])
+        tree = KDTree(infected, leaf_size = 20)
+        # We query the neighbors based on the radius of infection
+        indices = tree.query_radius(node, r=self.radius)
+
+        if(len(indices[0]) > 0):
+             # Added another layer of infection checking by using the infection rate
+             state = np.random.choice( 
+                    [False, True], 
+                    1,
+                    p=[1-self.infectionRate, self.infectionRate]
+                    )
+             return state[0]
+        else:
+            return False
+
+    def updateNodes(self):
+        """
+        Updates all node position and updates new infections, recoveries and deaths.
+        """
+        infections = []
+        recoveries = []
+        deaths = []
+
+        for node in self.nodes_all["Susceptible"]:
+            node.takeStep()
+            if node.state == "Susceptible": 
+                infected = self.checkInfection(node)
+                if infected==True:
+                    infections.append(node)
+
+        for node  in self.nodes_all['Infected']:
+            node.takeStep() 
+            if node.state == "Recovered":
+                recoveries.append(node)  
+            if node.state == "Dead":
+                deaths.append(node)
+
+        for node in self.nodes_all['Recovered']:
+            node.takeStep()
+
+        for infection in infections:
+            infection.changeState("Infected")
+            self.nodes_all["Susceptible"].remove(infection)
+            self.nodes_all["Infected"].append(infection)
+
+        for recovery in recoveries: 
+             self.nodes_all["Infected"].remove(recovery)
+             self.nodes_all["Recovered"].append(recovery)
+
+        for death in deaths: 
+             self.nodes_all["Infected"].remove(death)
+             self.nodes_all["Dead"].append(death)
+
+    def getInfectedCoordinates(self):
+        """
+        Gets the coordinates of all infected nodes. Coordinates are represented
+        in a tuple.
+
+        Returns
+        -------
+        nodes : list
+            List of 2-tuple coordinate vectors.
+
+        """
+        nodes = self.nodes_all['Infected']
+        nodes = map(lambda node: node.getCoordinates(), nodes)
+        nodes = list(nodes)
+        return nodes 
+
+    def getCoordinates(self, name):
+        """
+        Gets the coordinates of nodes. Coordinates are represented
+        in a tuple.
+
+        Returns
+        -------
+        nodes : list
+            List of 2-tuple coordinate vectors.
+
+        """
+        nodes = self.nodes_all[name]
+        nodes = map(lambda node: node.getCoordinates(), nodes)
+        nodes = list(nodes)
+        return nodes 
+
+    def getSample(self, name):
+        """
+        Parameters
+        ----------
+        name : String
+            State of nodes desired.
+
+        Returns
+        -------
+        Two lists of node coordinates.
+        """
+        nodes = self.nodes_all[name]
+        nodes = map(lambda node: node.getCoordinates(), nodes)
+        nodes = list(nodes)
+
+        X = map(lambda node: node[0], nodes)
+        Y = map(lambda node: node[1], nodes)
+
+        X = list(X)
+        Y = list(Y)
+        return X,Y
+
+    def getApiResponse(self):
+        """
+        I wrote this function so the coordinates, and colors of the nodes can 
+        be exported as an API response to be used in different visualization
+        frameworks. 
+
+        Returns
+        -------
+        nodes : list of dictionaries
+            Each dict object contains the coordinates and color of the node.
+
+        """
+        states = ["Infected", "Susceptible", "Recovered", "Dead"]
+        nodes =[]
+        for state in states:
+            nodes += self.nodes_all[state]      
+        func = lambda node: {"Coord": node.getCoordinates(), "Color": node.getColor()}
+
+        nodes = map(func, nodes)
+        nodes = list(nodes)
+        return nodes
+
+    def getInfectionNo(self):
+        return len(self.nodes_all['Infected'])
+
+    def getSusceptibleNo(self):
+      return len(self.nodes_all['Susceptible'])
+
+    def getRecoveredNo(self):
+      return len(self.nodes_all['Recovered'])
+
+    def getDeathsnNo(self):
+      return len(self.nodes_all['Dead'])
+
+
+
+class Node(object):
+    def __init__(self, max_coordinate, state, speed):
+        self.x = r.randint(-max_coordinate*2,max_coordinate*2)
+        self.y = r.randint(-max_coordinate,max_coordinate)       
+        self.state = state
+        self.colorMap = {
+            "Susceptible": "blue",
+            "Infected": "red",
+            "Recovered": "Green",
+            "Dead": "black"}
+        self.counter = 0
+        self.angle = np.random.uniform(0,2*np.pi)
+        self.max = max_coordinate
+        self.speed = speed
+
+    def getCoordinates(self):
+        return [self.x,self.y] 
+
+    def takeStep(self):
+        if self.state == "Infected":
+            if self.counter == 400:
+
+                state = np.random.choice( 
+                    ['Dead', 'Recovered'], 
+                    1,
+                    p=[0.1, 0.9]
+                    )
+                self.state = state[0]
+            else:
+                self.counter+=1
+
+        x = self.x
+        y = self.y
+        y += np.cos(self.angle)*self.speed
+        x += np.sin(self.angle)*self.speed
+        if abs(x) < self.max*2:
+            self.x = x
+        else:
+            self.angle = np.random.uniform(0,2*np.pi)
+        if abs(y) < self.max - 5:
+            self.y = y
+        else:
+            self.angle = np.random.uniform(0,2*np.pi)
+
+    def state(self):
+        return self.state
+    def getColor(self):
+        return self.colorMap[self.state]
+    def changeState(self, state):
+        self.state = state
+
+    def __str__(self):
+        return str((self.x,self.y))
+
+    def __repr__(self):
+        return str((self.x,self.y))
+
+
+
+
+
+"""
+Below is a sample test run to show how the model is used. 
+"""
+
+
+# size = 300
+# nodes = 160
+# infectionRate = 0.5
+# radius = 7
+# speed = 1.9
+# steps = 50
+# pandemic = PandemicSpread(size, nodes, infectionRate, radius, speed)
+# pandemic.startInfection()
+# pandemic.deployNodes()
+
+# for step in range(steps):
+#     pandemic.takeStep()
+
+
+
+
+
+

requirements.txt

@@ -0,0 +1,4 @@
+numpy>=1.18.1
+scipy>=1.4.1
+scikit-learn>=0.22.1
+matplotlib>=3.1.3