updated greedy algorithm

2 files changed, 18 insertions, 18 deletions

diff --git a/jpa_test/algorithms.py b/jpa_test/algorithms.py
index facce3b..cae16b8 100644
--- a/jpa_test/algorithms.py
+++ b/jpa_test/algorithms.py
@@ -10,28 +10,25 @@ import timeout
 def greedy_prediction(G, cascades):
     """
     Returns estimated graph from Greedy algorithm in "Learning Epidemic ..."
+    Only words for independent cascade model!
     """
-    #TODO: This function is deprecated!
     G_hat = cascade_creation.InfluenceGraph(max_proba=None)
     G_hat.add_nodes_from(G.nodes())
     for node in G_hat.nodes():
-        unaccounted = np.ones(len(cascades), dtype=bool)
-        for t, cascade in izip(xrange(len(cascades)), cascades):
-            if not cascade.infection_time(node) or \
-                    cascade.infection_time(node)[0] == 0:
-                unaccounted[t] = False
-        while unaccounted.any():
-            tmp = [cascade for boolean, cascade in izip(unaccounted,
-                    cascades) if boolean]
+        print node
+        # Avoid cases where infection time is None or 0
+        tmp = [cascade for cascade in cascades if cascade.infection_time(node)
+                                                                           [0]]
+        while tmp:
             parents = Counter()
             for cascade in tmp:
                     parents += cascade.candidate_infectors(node)
             parent = parents.most_common(1)[0][0]
             G_hat.add_edge(parent, node)
-            for t, cascade in izip(xrange(len(cascades)), cascades):
-                if (cascade.infection_time(parent) == \
-                        [item - 1 for item in cascade.infection_time(node)]):
-                    unaccounted[t] = False
+            tmp = [cascade for cascade in tmp if (
+                    cascade.infection_time(parent)[0] is not None and
+                    cascade.infection_time(parent)[0]+1 not in
+                        cascade.infection_time(node))]
     return G_hat
 
 
@@ -79,10 +76,13 @@ def test():
     G.erdos_init(n = 100, p = .05)
     import time
     t0 = time.time()
-    A = cascade_creation.generate_cascades(G, .2, 1000)
-    G_hat = recovery_l1obj_l2constraint(G, A,
-            passed_function=convex_optimization.l1obj_l2penalization,
-            floor_cstt=.1, lbda=10)
+    A = cascade_creation.generate_cascades(G, .2, 10000)
+    if 1:
+        G_hat = greedy_prediction(G, A)
+    if 0:
+        G_hat = recovery_l1obj_l2constraint(G, A,
+                passed_function=convex_optimization.l1obj_l2penalization,
+                floor_cstt=.1, lbda=10)
     correctness_measure(G, G_hat, print_values=True)
 
 
diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
index dbb9e54..9a26c03 100644
--- a/jpa_test/cascade_creation.py
+++ b/jpa_test/cascade_creation.py
@@ -65,7 +65,7 @@ class Cascade(list):
         """
         candidate_infectors = collections.Counter()
         for t in self.infection_time(node):
-            if t > 0:
+            if t:  # avoid cases where t=0 or t is None
                 candidate_infectors.update(np.where(self[t-1])[0])
         return candidate_infectors