From b8e67ac9ae3ea0dff7f3e3d69fa5a5bbb54f08ee Mon Sep 17 00:00:00 2001
From: Chris Vanderloo <chris.vanderloo@yahoo.com>
Date: Mon, 27 Apr 2020 20:25:03 -0400
Subject: [PATCH] small performance improvements

---
 actors.py      |  4 ++--
 epidemic.py    | 36 ++++++++++++++++--------------------
 interaction.py |  3 ++-
 3 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/actors.py b/actors.py
index 3ccc29d..5092246 100644
--- a/actors.py
+++ b/actors.py
@@ -355,8 +355,8 @@ def generate_synthetic(n):
             nhts_hh_templates), folder='synthetic_hh')
 
     print("Assigning activity locations.")
-    # 8 people per location (work, home, etc)
-    assign_dummy_locations(synthetic_households, int(n / 8))
+    # 4 people to a location avg (work, home, etc)
+    assign_dummy_locations(synthetic_households, int(n / 4))
     # assign_locations(synthetic_households)
 
     return synthetic_households
diff --git a/epidemic.py b/epidemic.py
index 10bcd1c..5aea9a6 100644
--- a/epidemic.py
+++ b/epidemic.py
@@ -147,12 +147,14 @@ class EpidemicSim:
             }})
 
     def increment_time(self, node):
-        n = self.get_attr(node, 'time_infected')
-        self.set_attr(node, 'time_infected', n + 1)
+        self.G.node[node]['time_infected'] += 1
 
     def get_attr(self, node, attr):
         return self.G.node[node][attr]
 
+    def get_attrs(self, node):
+        return self.G.node[node]
+
     def set_attr(self, node, attr, val):
         self.G.node[node][attr] = val
 
@@ -174,27 +176,25 @@ class EpidemicSim:
         acttype can be one of:
         (H)ome, (W)ork, (S)hop, s(C)hool, and (O)ther
         '''
-
         for n in self.get_people():
+            n_attrs = self.get_attrs(n)
             state = self.get_state(n)
             if state is 'E' or state is 'I' or state is 'Q':
                 self.increment_time(n)
-                if self.get_attr(n, 'infection_length') <= self.get_attr(
-                        n, 'time_infected'):
-                    if self.get_attr(n, 'will_die'):
+                if n_attrs['infection_length'] <= n_attrs['time_infected']:
+                    if n_attrs['will_die']:
                         self.update_state(n, 'D')
                     else:
                         self.update_state(n, 'R')
             if state == 'E' and \
-                    self.get_attr(n, 'incubation') == self.get_attr(n, 'time_infected'):
+                    n_attrs['incubation'] == n_attrs['time_infected']:
                 self.update_state(n, 'I')
-            if state == 'I':  # determine whether they should quarantine
-                if self.get_attr(n, 'test_submitted'):
+            elif state == 'I':  # determine whether they should quarantine
+                if n_attrs['test_submitted']:
                     self.set_attr(n, 'days_since_submitted_test',
-                                  self.get_attr(n, 'days_since_submitted_test') + 1
-                                  )
-                    if self.get_attr(n, 'days_since_submitted_test') == self.get_attr(
-                            n, 'test_turnaround'):
+                                    n_attrs['days_since_submitted_test'] + 1
+                                 )
+                    if n_attrs['days_since_submitted_test'] == n_attrs['test_turnaround']:
                         self.update_state(n, 'Q')
                         self.confirmed += 1
                 elif random.random() < self.config['test_rate']:
@@ -205,13 +205,8 @@ class EpidemicSim:
             u_state = self.get_state(u)
             v_state = self.get_state(v)
 
-            # if they've recovered, quarantined, or dead we ignore the
-            # interaction
-            if u_state == 'Q' or v_state == 'Q':
-                continue
-            if u_state == 'R' or v_state == 'R':
-                continue
-            if u_state == 'D' or v_state == 'D':
+            # if one is susceptible, then we may want to infect
+            if u_state != 'S' and v_state != 'S':
                 continue
 
             if random.random() < self.get_infection_on_interaction():
@@ -222,6 +217,7 @@ class EpidemicSim:
 
     def run_full_simulation(self, days, totalPeople):
         # Sort edges of graph by timestep
+
         potential_interactions = generate_interactions(self.G)
         finished = False
         infected = []
diff --git a/interaction.py b/interaction.py
index abd1bc4..0aa219a 100644
--- a/interaction.py
+++ b/interaction.py
@@ -1,5 +1,6 @@
 import networkx as nx
 import random
+from tqdm import tqdm
 
 '''
 Responsible for converting the environment interaction graph 
@@ -39,7 +40,7 @@ Where G is an generated graph from epidemic.generate_graph
 '''
 def generate_interactions(G):
   interactions = []
-  for person1 in G.nodes():
+  for person1 in tqdm(G.nodes()):
     if (str(person1).startswith('P_')):
       locations = [n for n in G.neighbors(person1)]
       for loc in locations:
-- 
2.51.2