bugfix

GiulioRossetti · Aug 9, 2017 · ae5d0b8 · ae5d0b8
1 parent 0bebf4e
commit ae5d0b8
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/ndlib/models/epidemics/SEIRModel.py b/ndlib/models/epidemics/SEIRModel.py
@@ -39,6 +39,8 @@ def __init__(self, graph):
             "edges": {},
         }
 
+        self.progress = {}
+
     def iteration(self, node_status=True):
         self.clean_initial_status(self.available_statuses.values())
 
@@ -66,8 +68,15 @@ def iteration(self, node_status=True):
                 infected_neighbors = len([v for v in neighbors if self.status[v] == 2 or self.status[v] == 1])
                 if eventp < self.params['model']['beta'] * infected_neighbors:
                     actual_status[u] = 2  # Exposed
-            elif 1 < u_status <= 2:
-                actual_status[u] -= self.params['model']['alpha']
+                    self.progress[u] = 0
+
+            elif u_status == 2:
+                if self.progress[u] < 1:
+                    self.progress[u] -= self.params['model']['alpha']
+                else:
+                    actual_status[u] = 1  # Infected
+                    del self.progress[u]
+
             elif u_status == 1:
                 if eventp < self.params['model']['gamma']:
                     actual_status[u] = 3  # Removed

diff --git a/ndlib/models/epidemics/SEISModel.py b/ndlib/models/epidemics/SEISModel.py
@@ -48,6 +48,7 @@ def __init__(self, graph):
         }
 
         self.name = "SEIS"
+        self.progress = {}
 
     def iteration(self, node_status=True):
         """
@@ -81,8 +82,14 @@ def iteration(self, node_status=True):
                 infected_neighbors = len([v for v in neighbors if self.status[v] == 2 or self.status[v] == 1])
                 if eventp < self.params['model']['beta'] * infected_neighbors:
                     actual_status[u] = 2  # Exposed
-            elif 1 < u_status <= 2:
-                    actual_status[u] -= self.params['model']['alpha']
+                    self.progress[u] = 0
+
+            elif u_status == 2:
+                if self.progress[u] < 1:
+                    self.progress[u] -= self.params['model']['alpha']
+                else:
+                    actual_status[u] = 1  # Infected
+                    del self.progress[u]
 
             elif u_status == 1:
                 if eventp < self.params['model']['lambda']: