1 files changed, 18 insertions, 6 deletions
diff --git a/jpa_test/algorithms.py b/jpa_test/algorithms.py
index 2ed015a..f3bf917 100644
--- a/jpa_test/algorithms.py
+++ b/jpa_test/algorithms.py
@@ -4,6 +4,7 @@ import cascade_creation
 from collections import Counter
 from itertools import izip
 import convex_optimization
+import timeout
 
 
 def greedy_prediction(G, cascades):
@@ -33,7 +34,7 @@ def greedy_prediction(G, cascades):
                     unaccounted[t] = False
     return G_hat
 
-
+@timeout.timeout(10)
 def recovery_l1obj_l2constraint(G, cascades):
     """
     Returns estimated graph from following convex program:
@@ -43,11 +44,15 @@ def recovery_l1obj_l2constraint(G, cascades):
     G_hat = cascade_creation.InfluenceGraph(max_proba=None)
     G_hat.add_nodes_from(G.nodes())
     for node in G_hat.nodes():
+        print node
         M, w = cascade_creation.icc_matrixvector_for_node(cascades, node)
         p_node, __ = convex_optimization.l1obj_l2constraint(M,w)
+        G_hat = cascade_creation.add_edges_from_proba_vector(G=G_hat,
+                    p_node=p_node, node=node, floor_cstt=.01)
+    return G_hat
 
 
-def correctness_measure(G, G_hat):
+def correctness_measure(G, G_hat, print_values=False):
     """
     Measures correctness of estimated graph G_hat to ground truth G
     """
@@ -56,6 +61,10 @@ def correctness_measure(G, G_hat):
     fp = edges_hat - edges
     fn = edges - edges_hat
     gp = edges | edges_hat
+    if print_values:
+        print "False Positives: {}".format(len(fp))
+        print "False Negatives: {}".format(len(fn))
+        print "Good Positives: {}".format(len(gp))
     return fp, fn, gp
 
 
@@ -63,13 +72,16 @@ def test():
     """
     unit test
     """
-    G = cascade_creation.InfluenceGraph(max_proba = .3)
-    G.erdos_init(n = 100, p = .5)
+    G = cascade_creation.InfluenceGraph(max_proba = .5)
+    G.erdos_init(n = 100, p = .3)
     import time
     t0 = time.time()
-    A = cascade_creation.generate_cascades(G, .2, 2)
+    A = cascade_creation.generate_cascades(G, .2, 50)
+
+    G_hat = recovery_l1obj_l2constraint(G, A)
+
+    correctness_measure(G, G_hat, print_values=True)
 
-    recovery_l1obj_l2constraint(G, A)
 
 if __name__=="__main__":
     test()